diff --git a/bin/main.ml b/bin/main.ml
index 3b84ffc..662c879 100644
--- a/bin/main.ml
+++ b/bin/main.ml
@@ -1,92 +1,144 @@
 (*let debugf = Format.ifprintf Format.std_formatter*)
 
-let pp_card out ((x, y), n) = Format.fprintf out "((%d, %d), %d)" x y n
+type pos = int * int
+(** A position on the grid, pair of x, y co-ordinates *)
 
-let intersects ((x1l, y1b), n1) ((x2l, y2b), n2) =
-  let x1r = x1l + n1 in
-  let x2r = x2l + n2 in
-  let y1t = y1b + n1 in
-  let y2t = y2b + n2 in
-  let result = x1l < x2r && x1r > x2l && y1t > y2b && y1b < y2t in
-  begin
-    result
-  end
+(** Pretty print a position to Format.formatter [out] *)
+let pp_pos out ((x, y) : pos) = Format.fprintf out "(%d,@ %d)" x y
 
-let rec card_fits size idx placed_cards (x, y) =
-  if x + idx > size then false
-  else if y + idx > size then false
-  else
-    match placed_cards with
+(** 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 ((x, y), idx) then false else card_fits size idx t (x, y)
+    | 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
 
-let rec in_card (x, y) cards =
-  match cards with
+(** 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
-  | ((a, b), n) :: t ->
-      if x >= a && x < a + n && y >= b && y < b + n then true
-      else in_card (x, y) t
+  | 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
 
-let next_pos size (x, y) cards =
+(** 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 >= size then impl 0 (y + 1)
-    else if in_card (x, y) cards then impl (x + 1) 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 pp_pos out (x, y) = Format.fprintf out "(@[%d,@ %d@])" x y*)
+let triangle_num n = n * (n + 1) / 2
 
-let rec find_solutions_impl cards size n idx current_alloc current_pos =
-  begin
-    (*debugf "find_solutions_impl:@ @[<hov>%a@ %d@ %d@ %d@ %a@ %a@]@;@?"
-	(Format.pp_print_array Format.pp_print_int) cards
-	size n idx
-	(Format.pp_print_list pp_card) current_alloc
-	pp_pos current_pos;*)
-    if current_pos = (0, size) then current_alloc
+(** 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 cards.(idx) = 0 then
-      find_solutions_impl cards size n (idx - 1) current_alloc current_pos
-    else if card_fits size idx current_alloc current_pos then begin
-      Array.set cards idx (cards.(idx) - 1);
-      let new_alloc = (current_pos, idx) :: current_alloc in
-      let new_pos = next_pos size current_pos new_alloc in
-      let alloc = find_solutions_impl cards size n n new_alloc new_pos in
-      Array.set cards idx (cards.(idx) + 1);
-      if List.is_empty alloc then
-        find_solutions_impl cards size n (idx - 1) current_alloc current_pos
-      else alloc
+    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 find_solutions_impl cards size n (idx - 1) current_alloc current_pos
-  end
-
-let find_solutions cards size =
-  find_solutions_impl cards size
-    (Array.length cards - 1)
-    (Array.length cards - 1)
-    [] (0, 0)
-
-exception Overlapping_value
-
-let print_solution size cards =
-  let array = Array.make (size * size) '.' in
-  let set_pos x y c =
-    if array.(x + (y * size)) <> '.' then raise Overlapping_value
-    else Array.set array (x + (y * size)) c
+    else impl (idx - 1) pos sqs
   in
-  let rec write_size x y n =
+  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 * size)) (Char.chr (48 + (n mod 10)));
-      write_size (x - 1) y (n / 10)
+      Array.set array (x + (y * length)) (Char.chr (48 + (n mod 10)));
+      write_length (x - 1) y (n / 10)
     end
   in
-  let rec impl cards =
-    match cards with
+  let rec impl sqs =
+    match sqs with
     | [] -> ()
-    | ((x, y), n) :: t -> begin
+    | { 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 '+';
@@ -108,28 +160,23 @@ let print_solution size cards =
             set_pos x (y + a) '|';
             set_pos (x + n - 1) (y + a) '|'
           done;
-          write_size (x + n - 2) (y + 1) n
+          write_length (x + n - 2) (y + 1) n
         end;
         impl t
       end
   in
-  impl cards;
-  for y = 0 to size - 1 do
-    for x = 0 to size - 1 do
-      Format.printf "%c" array.(x + (y * size))
+  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 = (n + 1) * n / 2
+let tri_n = triangle_num n
 
-(* Cards is an array initialised so that cards[x] = x for x = [0..9].
-
-   These are the cards we need to fit into the square of length tri_n.
- *)
-let cards = Array.init (n + 1) Fun.id
-let soln = find_solutions cards tri_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_card) soln
+let () = Format.printf "Solution: %a@]@\n" (Format.pp_print_list pp_square) soln
 let () = print_solution tri_n soln