(* Lecture 11: N-Queens *)

(* Workaround for Harlequin "restricted beta" problem *)
(* This make 'a option, SOME, NONE available at the top-level *)
open General;

(* Signature for dictionaries *)
(*

   Unlike the code in class, here we use pairs of integers as
   keys.  They are compared with the function inPairCompare.

   Existing entries may be "updated" by inserting a new entry
   with the same key.
*)

fun intPairCompare ((i1,i2), (j1,j2)) =
  (case Int.compare (i1,j1)
     of EQUAL => Int.compare (i2,j2)
      | GREATER => GREATER
      | LESS => LESS);

signature DICT =
sig
  type key = int * int
  type 'a entry = key * 'a

  type 'a dict

  val empty : 'a dict
  val lookup : 'a dict -> key -> 'a option
  val insert : 'a dict * 'a entry -> 'a dict
  val mapEntries : 'a dict -> ('a entry -> 'a) -> 'a dict
end;  (* signature DICT *)

(* Straight binary search tree implementation of dictionaries *)

structure BinarySearchTree :> DICT =
struct
  type key = int * int
  type 'a entry = key * 'a

  datatype 'a tree =
    Empty
  | Node of 'a entry * 'a tree * 'a tree

  type 'a dict = 'a tree

  (* Representation Invariant *)
  (*
     The tree is ordered:
       For every node  Node((key1,datum1), left, right),
       every key in left is LESS than key1, and
       every key in right is GREATER than key1
  *)

  val empty = Empty

  fun lookup tree key =
    let fun lk (Empty) = NONE
	  | lk (Node((key1,datum1), left, right)) =
              (case intPairCompare(key,key1)
		 of EQUAL => SOME(datum1)
	          | LESS => lk left
	          | GREATER => lk right)
    in
      lk tree
    end

  fun insert (tree, entry as (key,datum)) =
    (* val ins : 'a dict -> 'a dict *)
    (* ins (tree) inserts entry into tree *)
    let fun ins (Empty) = Node(entry, Empty, Empty)
	  | ins (Node(entry1 as (key1,datum1), left, right)) =
              (case intPairCompare(key,key1)
		 of EQUAL => Node(entry, left, right)
	          | LESS => Node(entry1, ins left, right)
                  | GREATER => Node(entry1, left, ins right))
    in 
      ins tree
    end

  fun mapEntries (Empty) f = Empty
    | mapEntries (Node(e as (key,datum),left,right)) f =
         Node ((key, f(e)), mapEntries left f,
	       mapEntries right f)
end;  (* structure BinarySearchTree *)

structure Dict :> DICT = BinarySearchTree;

datatype square = Empty | Queen | Blocked;

(* initialize n b (i,j) where n > 0, 1 <= i,j <= n *)
(* creates empty squares between i and j and (1,1) *)
fun initialize n b (i,j) =
       next n (Dict.insert (b,((i,j),Empty))) (i,j)
and next n b (1,1) = b
  | next n b (i,1) = initialize n b (i-1,n)
  | next n b (i,j) = initialize n b (i,j-1);

(* emptyBoard n creates n * n empty board *)
fun emptyBoard n = initialize n (Dict.empty) (n,n);

(*
   block (i,j) ((x,y),Empty)
    >=> Blocked if (i,j) attacks (x,y),
    >=> Empty otherwise
   Other squares remain untouched
 *)
fun block (i:int,j:int) (e as ((x,y),Empty)) =
       if i = x (* same column *)
	  orelse y = j (* same row *)
	  orelse x-i = y-j (* same positive diagonal *)
	  orelse x-i = j-y (* same negative diagonal *)
       then Blocked
       else Empty
  | block (i,j) (_,sq) = sq;  (* Queen or Blocked *)

(* place n b (i,j)  places Queen on (i,j) in b and continues *)
fun place n b (i,j) =
    let
      val b' = Dict.mapEntries b (block (i,j))
    in
      nextColumn n (Dict.insert (b',((i,j),Queen))) (i-1,j)
    end
(* nextColumn n b (i,j) places on column i or succeeds *)
and nextColumn n b (0,j) = true
  | nextColumn n b (i,j) = nextRow n b (i,n)
(* nextrow n b (i,j) tries to place on row below j or fails *)
and nextRow n b (i,0) = false
  | nextRow n b (i,j) = check n b (Dict.lookup b (i,j)) (i,j)
(* check n b (i,j) places on (i,j) if OK or trues next row *)
and check n b (SOME(Empty)) (i,j) =
       place n b (i,j) orelse nextRow n b (i,j-1)
  | check n b _ (i,j) = nextRow n b (i,j-1);  (* Blocked *)
    (* NONE or Queen should be impossible *)

(* solvable n Checks if n-Queens has solution for n*n board *)
fun solvable n = check n (emptyBoard n) (SOME(Empty)) (n,n);