(*let debugf = Format.ifprintf Format.std_formatter*)

type pos = int * int
(** A position on the grid, pair of x, y co-ordinates *)

(** Pretty print a position to Format.formatter [out] *)
let pp_pos out ((x, y) : pos) = Format.fprintf out "(%d,@ %d)" x y

(** Get the x co-ordinate of a position *)
let pos_x = fst

(** Get the y co-ordinate of a position *)
let pos_y = snd

type square = { pos : pos; length : int }
(** A type representing a square, consisitng of the bottom-left corner of the
    square and the length of each side. *)

(** Pretty print a square to Format.formatter [out] *)
let pp_square out (sq : square) =
  Format.fprintf out "{%a@ len:%d}" pp_pos sq.pos sq.length

(** Returns true if the squares [sq1] and [sq2] intersect, and false otherwise.
*)
let intersects sq1 sq2 =
  let sq1l = pos_x sq1.pos in
  let sq1r = sq1.length + sq1l in
  let sq1b = pos_y sq1.pos in
  let sq1t = sq1.length + sq1b in
  let sq2l = pos_x sq2.pos in
  let sq2r = sq2.length + sq2l in
  let sq2b = pos_y sq2.pos in
  let sq2t = sq2.length + sq2b in
  sq1l < sq2r && sq1r > sq2l && sq1t > sq2b && sq1b < sq2t

(** Returns true if we can place the square [sq] without overlapping any already
    placed squares in [sqs] and without exceeding the bounds of the grid which
    is [length] along each side. *)
let square_fits sq length sqs =
  let rec impl sq sqs =
    match sqs with
    | [] -> true
    | h :: t -> if intersects h sq then false else impl sq t
  in
  if pos_x sq.pos + sq.length > length then false
  else if pos_y sq.pos + sq.length > length then false
  else impl sq sqs

(** Returns true if the position [(x, y)] is in one of the squares [sqs]. *)
let rec in_squares ((x, y) : pos) sqs =
  match sqs with
  | [] -> false
  | h :: t ->
      let sqx = pos_x h.pos in
      let sqy = pos_y h.pos in
      let len = h.length in
      if x >= sqx && x < sqx + len && y >= sqy && y < sqy + len then true
      else in_squares (x, y) t

(** Returns the next position to consider when working through the grid we are
    placing squares on. [(x, y)] is the current position, and [sqs] is a list of
    already placed squares.

    Returns (0, length) when we have filled the grid. *)
let next_pos length ((x, y) : pos) sqs =
  (* We basically walk along each row looking for an empty space. *)
  let rec impl x y =
    if x >= length then impl 0 (y + 1)
    else if in_squares (x, y) sqs then impl (x + 1) y
    else (x, y)
  in
  impl (x + 1) y

let triangle_num n = n * (n + 1) / 2

(** Find a solution to the [n]th Partridge problem. Returns a list of squares
    giving the position on the grid. *)
let find_solution n =
  (* recursive implementation:

	   [avail_sqs] is an array where avail_sqs.(x) returns how many sqs of size 
	   x are available to be placed.  [length] is the side length of the grid
	   we are placing the squares into.

	   [impl idx pos sqs] implements the recursive implementation.  [idx] is the
	   current index in [avail_sqs] that we are looking at.

	   If there is a square available of size [idx] (i.e. avail_sqs.(idx) > 0) 
	   then we try to place a square of size [idx] at [pos].  If this is
	   successful it adds that to the list [sqs] and tries to find a square that
	   fits in the next position.

	   If [impl] is not successful it tries again at the current position with
	   a square of size [idx - 1].

	   If we reach an [idx] of 0 we have failed and return an empty list.

	   If we reach the position [(0, length)] we have succeeded and return [sqs].

	*)
  let avail_sqs = Array.init (n + 1) Fun.id in
  let length = triangle_num n in
  let rec impl idx pos sqs =
    let sq = { pos; length = idx } in
    if pos = (0, length) then sqs
    else if idx = 0 then []
    else if avail_sqs.(idx) = 0 then impl (idx - 1) pos sqs
    else if square_fits sq length sqs then begin
      Array.set avail_sqs idx (avail_sqs.(idx) - 1);
      let new_sqs = sq :: sqs in
      let new_pos = next_pos length pos new_sqs in
      let result = impl n new_pos new_sqs in
      Array.set avail_sqs idx (avail_sqs.(idx) + 1);
      if List.is_empty result then impl (idx - 1) pos sqs else result
    end
    else impl (idx - 1) pos sqs
  in
  impl n (0, 0) []

(** Exception raised if we find we have overlapping squares when printing the
    solution. *)
exception Overlapping_squares of pos

(** Print the layout given in [sqs] for a grid with side-length [size]. *)
let print_solution length sqs =
  let array = Array.make (length * length) '.' in
  let set_pos x y c =
    if array.(x + (y * length)) <> '.' then raise (Overlapping_squares (x, y))
    else Array.set array (x + (y * length)) c
  in
  let rec write_length x y n =
    if n = 0 then ()
    else begin
      Array.set array (x + (y * length)) (Char.chr (48 + (n mod 10)));
      write_length (x - 1) y (n / 10)
    end
  in
  let rec impl sqs =
    match sqs with
    | [] -> ()
    | { pos = x, y; length = n } :: t -> begin
        if n = 1 then set_pos x y '*'
        else if n = 2 then begin
          set_pos x y '+';
          set_pos (x + 1) y '+';
          set_pos x (y + 1) '+';
          set_pos (x + 1) (y + 1) '+'
        end
        else begin
          set_pos x y '+';
          set_pos (x + n - 1) y '+';
          set_pos x (y + n - 1) '+';
          set_pos (x + n - 1) (y + n - 1) '+';
          for a = 1 to n - 2 do
            set_pos (x + a) y '-';
            set_pos (x + a) (y + n - 1) '-';
            for b = 1 to n - 2 do
              set_pos (x + a) (y + b) ' '
            done;
            set_pos x (y + a) '|';
            set_pos (x + n - 1) (y + a) '|'
          done;
          write_length (x + n - 2) (y + 1) n
        end;
        impl t
      end
  in
  impl sqs;
  for y = 0 to length - 1 do
    for x = 0 to length - 1 do
      Format.printf "%c" array.(x + (y * length))
    done;
    Format.printf "\n"
  done

let n = 8
let tri_n = triangle_num n

let soln = find_solution n
let () = Format.printf "@[<hov>Base number: %d,@;side length: %d@;" n tri_n
let () = Format.printf "Solution: %a@]@\n" (Format.pp_print_list pp_square) soln
let () = print_solution tri_n soln