open Types
open Board
open Protocol

let danger () =
  let approx_dist (x1, y1) (x2, y2) = abs (x1 - x2) + abs (y1 - y2) in
  let my_loc0 = my_loc () in
  let robots =
    IntMap.fold
      (fun id robot acc ->
	if id != !my_id && approx_dist my_loc0 robot.robot_loc <= 4 then
	  robot.robot_loc :: acc
	else
	  acc)
      !robots
      [] in
  robots <> [] && (* approximate in advance (for efficiency) *)
  let dists, dtable =
    Dijkstra.dijkstra_list
      ~adjustment:Dijkstra.adjustment_simple
      ~d_limit:4.
      robots in
  (Dijkstra.earray dists my_loc0 <= 4.)

let rec send_or_fight proposed_command =
  let d, c = (Forecast.select_command proposed_command) in
  Schedule.aging ();
  send_cmd d c
    
let rec move (path, dest) =
  match path with
  | [] ->
      assert(my_loc () = dest);
      plan ()
  | next :: rest ->
      begin
	dirty := false;
	let my_loc0 = my_loc () in
	send_or_fight (Move next);
	if !dirty || my_loc () <> move_pos my_loc0 next then begin
	  Schedule.add_failed_location dest;
	  plan () 
	end
	else
	  move (rest, dest)
      end
and plan () =
  begin
    print_endline "delivering packages...";
    let my_loc1 = my_loc () in
    let delivers =
      IntSet.filter
	(fun d ->
	  match IntMap.find d !packs with
	  | Some info, No -> info.dest = my_loc1
	  | _ -> assert(1 = 2); true)
	(my_pack ()) in
    if not (IntSet.is_empty delivers) then
      send_or_fight (Drop (intset_to_list delivers));

    let my_loc2 = my_loc () in
    if my_loc2 <> my_loc1 then plan () else
    begin
      print_endline "clustering packages...";
      let clusters = Mst.find_mst () in
      (*List.iter
	(fun (w, s) ->
	  Format.printf "[DEBUG] cluster: w = %d, s = %s@." w
	    (Printer.print_int_list (intset_to_list s)))
	clusters;*)
      
      print_endline "selecting packages...";
      let dists, _ = Dijkstra.dijkstra my_loc2 in
      let clusters =
	List.map
	  (function w, s ->
	    match IntMap.find (IntSet.choose s) !packs with
	    | Some(info), loc ->
		assert(loc = No || loc = Yes(my_loc2));
		let d = Dijkstra.earray dists info.dest in
		float_of_int w, s, d
	    | _ -> assert(1 = 2); (0., IntSet.empty, max_float))
	  clusters in
      let clusters =
	List.sort
	  (fun (w1, _, d1) (w2, _, d2) -> compare (w2 *. d1) (w1 *. d2))
	  clusters in
      (*List.iter
	(fun (w, s, d) ->
	  Format.printf "[DEBUG] cluster (sorted): w = %d, s = %s, d = %f@."
	    (int_of_float w)
	    (Printer.print_int_list (intset_to_list s)) d)
	clusters;*)
      let rest = ref !max_weight in
      let picks =
	List.fold_left
	  (fun picks (_, s, _) ->
	    IntSet.fold
	      (fun id picks ->
		match IntMap.find id !packs with
		| None, _ -> assert(1 = 2); picks
		| Some(info), _ ->
		    if info.weight > !rest then picks else
		    (rest := !rest - info.weight;
		     id :: picks))
	      s
	      picks)
	  []
	  clusters in
      assert (!rest >= 0);
      (*Format.printf "[DEBUG] picks = %s@."
	(Printer.print_int_list picks);*)

      print_endline "dropping packages...";
      let drops =
	IntSet.filter (fun d -> not (List.mem d picks)) (my_pack ()) in
      if not (IntSet.is_empty drops) then
	send_or_fight (Drop (intset_to_list drops));

      let my_loc3 = my_loc () in
      if my_loc3 <> my_loc2 then plan () else
      begin
	print_endline "picking packages...";
	let picks =
	  let my_pack = my_pack () in
	  List.filter (fun d -> not (IntSet.mem d my_pack)) picks in
	(*Format.printf "[DEBUG] picks (real) = %s@."
	  (Printer.print_int_list picks);*)
	if picks <> [] then
	  send_or_fight (Pick picks);

	let my_loc4 = my_loc () in
	if my_loc4 <> my_loc3 then plan () else
        (* [TODO] If you fail to reach your destination too often,
	   you may give up the destination. *)
	print_endline "selecting a destination...";
	let dists, paths = Dijkstra.dijkstra my_loc4 in
	let dest = ref (0, 0) in
	let dist = ref max_float in
	List.iter
	  (fun id ->
	    match IntMap.find id !packs with
	    | Some(info), No ->
		let d = Dijkstra.earray dists info.dest *. 0.4 in
		let d = if Schedule.too_often info.dest then d *. 10. else d in
		if d < !dist then begin
		  dest := info.dest;
		  dist := d
		end
	    | _ -> assert(1 = 2))
	  (intset_to_list (my_pack ()));
	IntMap.iter
	  (fun _ (info, loc) ->
	    match info, loc with
	    | Some info, Yes(loc)
	      when my_weight () + info.weight <= !max_weight ->
		let d = Dijkstra.earray dists loc *. 1. in
		let d = if Schedule.too_often loc then d *. 10. else d in
		if d < !dist then begin
		  dest := loc;
		  dist := d
		end
	    | Some info, Maybe(loc)
	      when my_weight () + info.weight <= !max_weight ->
		let d = Dijkstra.earray dists loc *. 100. in
		let d = if Schedule.too_often loc then d *. 10. else d in
		if d < !dist then begin
		  dest := loc;
		  dist := d
		end
	    | None, Yes(loc) -> assert(1 = 2)
	    | None, Maybe(loc) ->
		let d = Dijkstra.earray dists loc *. 1000. in
		let d = if Schedule.too_often loc then d *. 10. else d in
		if d < !dist then begin
		  dest := loc;
		  dist := d
		end
	    | _ -> ())
	  !packs;
	LocSet.iter
	  (fun loc ->
	    match !board.(snd loc).(fst loc) with
	    | Home(n) ->
		let d = Dijkstra.earray dists loc *.
		    10. *. (1.1 ** float_of_int n) in
		let d = if Schedule.too_often loc then d *. 10. else d in
		if d < !dist then begin
		  dest := loc;
		  dist := d
		end
	    | _ -> assert(1 = 2))
	  !homes;
	if !dest = (0, 0) then begin
	  (* nowhere to go *) 
	  send_or_fight (Pick []);
	  plan ()
	end else begin
	  print_endline "computing a path...";
	  let path = Dijkstra.retrieve_path paths my_loc4 !dest in
	  print_endline "move, move, ...";
	  (*Printf.printf "[DEBUG] move: dist = %f, path = " !dist;*)
	  (*List.iter
	    (function
	      | N -> print_char 'N'
	      | E -> print_char 'E'
	      | S -> print_char 'S'
	      | W -> print_char 'W')
	    path;
	  print_newline ();*)
	  move (path, !dest)
	end
      end
    end
  end