open Types
open Board
open Map_analyze

type 'a elem_array = 
    'a array

let earray t (x, y) = t.(fastloc_encode x y)
let put_earray t (x, y) v = t.(fastloc_encode x y) <- v

let earray_old t (x, y) = t.(y).(x)
let earray_old3 t p = t.(fastloc_get_y p).(fastloc_get_x p)

type robot_table = No_robot | Robot | Adjacent

let robot_board = ref [||]

let make_robot_board () = 
(*Printf.printf "make_robot_board\n";*)
  let my_id = !my_id in
  let b = make_fastlocarray (!board_width + 2) (!board_height + 2) No_robot in
  robot_board := b;
  IntMap.iter
    (fun id desc ->
      if id <> my_id then begin
	let loc = fastloc_of_loc desc.robot_loc in
	b.(loc + 1) <- Adjacent;
	b.(loc - 1) <- Adjacent;
	b.(loc + fastloc_y_mult) <- Adjacent;
	b.(loc - fastloc_y_mult) <- Adjacent
      end)
    !robots;

  IntMap.iter
    (fun id desc ->
      if id <> my_id then begin
	let loc = fastloc_of_loc desc.robot_loc in
	b.(loc) <- Robot
      end) !robots
      
type elem = int * int

type heap = 
    { dtable : float array ;
      heappos : int array ;
      prev_dir : move array ;
      table : int array ;
      mutable size : int ;
      d_limit : float;
      stop_fun : fast_loc -> bool;
      mutable stop_limit : int;
    }

let dz = function
    S -> -fastloc_y_mult
  | N -> fastloc_y_mult
  | E -> 1
  | W -> -1

let swap h n1 n2 = 
  Array.set h.heappos h.table.(n1) n2;
  Array.set h.heappos h.table.(n2) n1;
  let tx = h.table.(n1) in
  h.table.(n1) <- h.table.(n2);
  h.table.(n2) <- tx

let rec downheap h node = 
  let l = node * 2 in
  if l > h.size then () else
  let target = 
    if l < h.size && 
      Array.get h.dtable (h.table.(l + 1)) < Array.get h.dtable (h.table.(l))
    then l + 1 else l
  in
  if Array.get h.dtable (h.table.(node))
      <= Array.get h.dtable (h.table.(target)) then ()
  else begin
    swap h node target;
    downheap h target
  end

let rec upheap h node = 
  if node <= 1 then () else 
  let target = node / 2 in
  if Array.get h.dtable (h.table.(target))
      <= Array.get h.dtable (h.table.(node)) then ()
  else begin
    swap h node target;
    upheap h target
  end

let infty = max_float

let rec dijkstra_loop ~adjustment_f h = 
  let p = h.table.(1) in
  if Array.get h.dtable p >= h.d_limit then
    ()
  else begin
    swap h 1 h.size;
    h.size <- h.size - 1;
    if h.size > 0 then downheap h 1;
    
    let prev_dir = Array.get h.prev_dir p in
    let prev_dist = Array.get h.dtable p in
    List.iter
      (fun 
	direction ->
	  let next = p + dz direction in
	  match !fast_board.(next) with
	    Plain | Home _ ->
	      let new_dist =
		prev_dist +.
		  adjustment_f prev_dist p next prev_dir direction
	      in
	      let heappos = Array.get h.heappos next in
	      if heappos <= h.size &&
		Array.get h.dtable next > new_dist then begin
		  Array.set h.dtable next new_dist;
		  Array.set h.prev_dir next direction;
		  upheap h heappos
		end
	  | Water | Wall -> ())
      [N; E; W; S];
    if h.stop_fun p then begin
(*      Printf.printf "stoploc found at %d,%d (dist %f, rest %d)\n"
	(fastloc_get_x p) (fastloc_get_y p) prev_dist (h.stop_limit-1);*)
      if h.stop_limit <= 1 then ( (*print_endline "search stopped."*) )
      else (h.stop_limit <- h.stop_limit - 1; dijkstra_loop ~adjustment_f h)
    end
    else dijkstra_loop ~adjustment_f h
  end

let adjustment_multiplay =
  (fun dist now next prev_dir now_dir ->
  let robot_weight =
    match !robot_board.(next) with
      No_robot -> 1.0
    | Robot -> max 5.0 (16.0 -. dist *. dist)
    | Adjacent -> max 1.25 (4.0 -. dist *. dist)
  in
  robot_weight *. !weight_map.(next) *. (if prev_dir = now_dir then 1.01 else 1.)
  ), make_robot_board

let adjustment_singleplay =
  (fun _ _ _ prev_dir now_dir -> if prev_dir = now_dir then 1.000001 else 1.),
  (fun () -> ())

let adjustment_simple =
  (fun _ _ _ _ _ -> 1.0),
  (fun () -> ())

let is_multiplay () = 
  try
    ignore
      (IntMap.fold
	 (fun _ _ issecond ->
	   if issecond then raise Exit else true) !robots false);
    false
  with
    Exit -> true

let get_adjustment_function () = 
      if is_multiplay () 
      then adjustment_multiplay
      else adjustment_singleplay

let timer = ref 0.0

let start_timer () = () (* timer := Unix.gettimeofday ()*)
let stop_timer s = () (* Printf.printf "%s: %g sec\n" s (Unix.gettimeofday () -. !timer); start_timer ()*)

let dijkstra_multi ?adjustment ?(d_limit = infty)  ?(stop_fun = fun _ -> false) ?(stop_limit = 1) starts = 
  let adjustment_f, adjustment_init = match adjustment with
    None -> get_adjustment_function ()
  | Some adj -> adj
  in
  (* initialize *)
  adjustment_init ();
(*  start_timer (); *)
  let dtable =
    Array.create
      ((!board_height+1) * fastloc_y_mult) infty
  in
(*  stop_timer "dtable_make";*)
  let prev_dir =
    Array.create
      ((!board_height+1) * fastloc_y_mult) N in
(*  stop_timer "prevdir_make";*)
  let heappos = 
    Array.create ((!board_height+1) * fastloc_y_mult) 0 in
  let heap = Array.make (!board_width * !board_height + 1) 0 in
(*  stop_timer "heappos_make";*)
  let heapsize = ref 0 in
  IntSet.iter 
    (fun p ->
      incr heapsize;
      heap.(!heapsize) <- p;
      dtable.(p) <- 0.0)
    starts;
  for y = 1 to !board_height do
    for x = 1 to !board_width do
      let p = fastloc_encode x y in
      if not (IntSet.mem p starts) then begin
	incr heapsize;
	heap.(!heapsize) <- p;
      end
    done
  done;
(*  stop_timer "heap_init2";*)
(*  Printf.printf "heapsize = %d\n" !heapsize;*)
  for t = 1 to !heapsize do
    Array.set heappos heap.(t) t;
  done;
(*  stop_timer "heappos_init";*)
  let h = { table = heap; 
	    heappos = heappos;
	    dtable = dtable;
	    prev_dir = prev_dir;
	    size = !heapsize;
	    d_limit = d_limit;
	    stop_fun = stop_fun;
	    stop_limit = stop_limit;
	  } in
  dijkstra_loop ~adjustment_f h;
(*  stop_timer "main loop";*)
  h.dtable, h.prev_dir

let dijkstra ?adjustment ?d_limit ?stop_fun ?stop_limit start = 
  dijkstra_multi ?adjustment ?d_limit ?stop_fun ?stop_limit
    (IntSet.singleton (fastloc_of_loc start))
    
let dijkstra_list ?adjustment ?d_limit ?stop_fun ?stop_limit starts = 
  dijkstra_multi ?adjustment ?d_limit ?stop_fun ?stop_limit
    (List.fold_left 
       (fun set loc -> IntSet.add (fastloc_of_loc loc) set)
       IntSet.empty starts)

let rec retrieve_path_fast acc path_info start dest =
  if dest = start then acc
  else begin
    let dir = path_info.(dest) in
    retrieve_path_fast (dir::acc) path_info start (dest - dz dir)
  end

let rec retrieve_path path_info start dest =
  retrieve_path_fast [] path_info (fastloc_of_loc start) (fastloc_of_loc dest)

let dijkstra_stopset ?adjustment from set =
  let n = IntSet.cardinal set in
  let found_p p = IntSet.mem p set in
  dijkstra ?adjustment ~stop_limit:n ~stop_fun:found_p from

let get_distances ?adjustment from to_list = 
  let ps = List.fold_left 
      (fun set elem -> IntSet.add (fastloc_of_loc elem) set)
      IntSet.empty to_list
  in
  let dist, path_info = dijkstra_stopset ?adjustment from ps in
  List.map 
    (fun p -> p, 
      earray dist p,
      retrieve_path path_info from p)
    to_list

let sort_by_distances ?adjustment from getpos elems = 
  let ps = List.fold_left
      (fun set elem -> IntSet.add (fastloc_of_loc (getpos elem)) set)
      IntSet.empty elems in
  let dist, path = dijkstra_stopset ?adjustment from ps in
  let l = List.map (fun o -> earray dist (getpos o), o) elems in
  List.sort compare l

let sort_packs_by_destination ?adjustment from ids = 
  sort_by_distances ?adjustment from 
    (fun id ->
      try
	match (IntMap.find id !packs) with
	  (Some { dest = p }, _) -> p
	| _ -> (0, 0)  (* infty *)
      with
	Not_found -> (0, 0))
    ids