module MazeParse = struct
  let split separator s =
    let list = ref [] in
    let start = ref 0 in
    let () = try
        while true do
          let index = String.index_from s !start separator in
          list := (String.sub s !start (index - !start)) :: !list;
          start := index + 1
        done
      with Not_found -> list := (String.sub s !start ((String.length s) - !start)) :: !list
    in
    List.rev !list

  let from_file filename =
    let channel = open_in_bin filename in
    let size = in_channel_length channel in
    let buffer = Buffer.create size in
    Buffer.add_channel buffer channel size;
    let s = Buffer.contents buffer in
    let lines = split '\n' s in
    let height, width = List.length lines, String.length (List.hd lines) in
    let maze = Array.make_matrix width height ' ' in
    for x = 0 to width -1 do
      for y = 0 to height -1 do
        maze.(x).(y) <- (List.nth lines y).[x]
      done
    done;
    maze
end

module Maze = struct
  let width maze = Array.length maze
  let height maze = Array.length maze.(0)
  let get maze (x, y) = maze.(x).(y)
  let set maze (x, y) value = maze.(x).(y) <- value

  let walk maze path =
    let walk_once coords = set maze coords '.' in
    List.iter walk_once path

  let draw maze =
    for y = 0 to (height maze) -1 do
      for x = 0 to (width maze) -1 do
        print_char (get maze (x, y))
      done;
      print_newline ()
    done

  let is_passable maze coord = get maze coord <> '#'

  let distance (x1, y1) (x2, y2) =
    let square x = x * x in
    let sqrt_int x = int_of_float (sqrt (float_of_int x)) in
    sqrt_int (square (x1 - x2) + square (y1 - y2))

  let neighbor_nodes maze (x, y) =
    let nodes = [] in
    let nodes = if x > 0 then (x -1, y) :: nodes else nodes in
    let nodes = if x < (width maze) - 1 then (x +1, y) :: nodes else nodes in
    let nodes = if y > 0 then (x, y -1) :: nodes else nodes in
    let nodes = if y < (height maze) - 1 then (x, y +1) :: nodes else nodes in
    nodes
end

exception No_path_found

type cost = { steps: int; estimation: int; camefrom: int * int }

type node =
  | Unknown
  | Closed
  | Exploring of cost
  | Explored of cost

let rec reconstruct_path matrix node path =
  let camefrom = match matrix.(fst node).(snd node) with
    | Explored cost -> cost.camefrom
    | Exploring cost -> cost.camefrom
    | _ -> node
  in
  if camefrom = node then node :: path
  else reconstruct_path matrix camefrom (node :: path)

let find_path maze start goal =
  let heuristic = Maze.distance in
  let matrix = Array.make_matrix (Maze.width maze) (Maze.height maze) Unknown in

  let module CoordSet = Set.Make(
    struct
      type t = int * int
      let estimate node = match node with
        | Exploring cost -> cost.estimation
        | Explored cost -> cost.estimation
        | Unknown -> max_int
        | Closed -> max_int

      let compare a b =
        let score_a = estimate matrix.(fst a).(snd a) in
        let score_b = estimate matrix.(fst b).(snd b) in
        let diff = score_a - score_b in
        if diff = 0 then Pervasives.compare a b else diff
    end) in

  let openset = ref (CoordSet.singleton start) in
  let select_best set = CoordSet.min_elt set in
  let is_explorable maze node =
    if not (Maze.is_passable maze node) then false
    else match matrix.(fst node).(snd node) with
      | Closed -> false
      | Unknown -> true
      | Exploring _ -> true
      | Explored _ -> false
  in
  let get_steps node = match matrix.(fst node).(snd node) with
    | Exploring cost -> cost.steps
    | Explored cost -> cost.steps
    | Closed -> max_int
    | Unknown -> failwith "should not be called"
  in
  let set_exploring node steps camefrom =
    matrix.(fst node).(snd node) <- Exploring { steps = steps; estimation = steps + (heuristic node goal); camefrom = camefrom };
  in
  let explored node = match matrix.(fst node).(snd node) with
    | Exploring cost -> Explored cost
    | _ -> failwith ("not really explored: " ^ string_of_int (fst node) ^ ", " ^ string_of_int (snd node))
  in
  set_exploring start 0 start;
  let rec loop () =
    if CoordSet.is_empty !openset then raise No_path_found
    else
      let current = select_best !openset in
      if current = goal then List.rev (reconstruct_path matrix goal [])
      else begin
        openset := CoordSet.remove current !openset;
        matrix.(fst current).(snd current) <- explored current;

        let neighbors = Maze.neighbor_nodes maze current in
        let explore node =
          if is_explorable maze node then
            let tentative_steps = (get_steps current) + 1 in
            if not (CoordSet.mem node !openset) then begin
              set_exploring node tentative_steps current;
              openset := CoordSet.add node !openset;
            end
            else if tentative_steps < (get_steps node) then set_exploring node tentative_steps current
        in
        List.iter explore neighbors;
        loop ()
      end
  in
  loop ()

let () =
  let maze = MazeParse.from_file "maze.txt" in
  let before = Unix.gettimeofday () in
  let path = find_path maze (0, 0) (78, 29) in

  for i = 0 to 1000 do
    ignore (find_path maze (0, 0) (78, 29))
  done;

  let after = Unix.gettimeofday () in
  Maze.walk maze path;
  Maze.draw maze;
  Printf.printf "%d steps\n" (List.length path);
  let time = (after -. before) *. 1000. in
  Printf.printf "%f ms\n" time