(* othello board implementation *)

type piece = Black | White
type move = Pass | Square of int * int


let nextSquare =
  let next = Array.create_matrix 8 8 (Array.create 8 None)
  and id x = x in
  let ops = [|pred, id; pred, succ; id, succ; succ, succ;
 	      succ, id; succ, pred; id, pred; pred, pred|] in
   (* need to make each copy the inner vector to each cell *)
  for row = 0 to 7 do
    for col = 0 to 7 do
      next.(row).(col) <- Array.create 8 None
    done
  done;
  for dir = 0 to 7 do
    let (drow, dcol) = ops.(dir) in
    for row = 0 to 7 do
      for col = 0 to 7 do
 	let nrow = drow row
 	and ncol = dcol col in
 	if nrow >= 0 && nrow < 8 && ncol >= 0 && ncol < 8 then
 	  next.(row).(col).(dir) <- Some(nrow, ncol)
      done
    done
  done;
  fun dir row col -> next.(row).(col).(dir)

(* given a row and column, returns a list of directions to search for possible
   moves in *)
let searchDirs =
  let tl = [2;3;4]			(*top left squares*)
  and t = [2;3;4;5;6]			(*top*)
  and tr = [4;5;6]			(*top right*)
  and l = [0;1;2;3;4]			(*left*)
  and r = [4;5;6;7;0]			(*right*)
  and c = [0;1;2;3;4;5;6;7]		(*center*)
  and bl = [0;1;2]			(*bottom left*)
  and b = [0;1;2;6;7]			(*bottom*)
  and br = [6;7;0] in			(*bottom right*)
  let dirs = [|[|tl; tl; t; t; t; t; tr; tr|];
	       [|tl; tl; t; t; t; t; tr; tr|];
	       [|l;  l;  c; c; c; c; r;  r|];
	       [|l;  l;  c; c; c; c; r;  r|];
	       [|l;  l;  c; c; c; c; r;  r|];
	       [|l;  l;  c; c; c; c; r;  r|];
	       [|bl; bl; b; b; b; b; br; br|];
	       [|bl; bl; b; b; b; b; br; br|]|] in
  fun r c -> dirs.(r).(c)

let adjSquares =
  let dirs = [0; 1; 2; 3; 4; 5; 6; 7] and
      adjTable = Array.create_matrix 8 8 [] in
  for row = 0 to 7 do
    for col = 0 to 7 do
      adjTable.(row).(col) <-
	 List.fold_left (fun lst dir ->
	   match (nextSquare dir row col) with
	     None -> lst
	   | Some(r, c) -> (r, c) :: lst)
	   [] dirs
    done
  done;
  fun r c -> adjTable.(r).(c)
	     


module type Board =
  sig
    type t
    val create : unit -> t
    val copy : t -> t
    val get : t -> int -> int -> piece option
    val set : t -> int -> int -> piece option -> unit
    (* board -> Board.piece -> direction -> row -> col -> (row, col) option *)
    (* direction is between 0 and 7, where 0 is north, 7 northwest. *)
    val sandwich : t -> piece -> int -> int -> int -> (int * int) option
    val legalMove : t -> piece -> move -> bool
    val makeMove : t -> piece -> move -> (t * int)
    val iter : (piece option -> 'a) -> t -> unit
    val iteri : (int -> int -> piece option -> 'a) -> t -> unit
  end

module ABoard : Board =
  struct
    type t = piece array array
	  
    let piece_of_num =
      let pieces = [|Some Black; Some White; None|] in
      function n -> pieces.(n)
    let num_of_piece = function
	Some Black -> 0
      |	Some White -> 1
      |	None -> 2
	    
    let create () = Array.create_matrix 8 8 None
    let copy = Array.map Array.copy
    let get brd row col = brd.(row).(col)
    let set brd row col pc = brd.(row).(col) <- pc

    let sandwich =
      let rec findEnd brd player dir r c =
    	match (nextSquare dir r c) with
	  Some(nr, nc) -> begin
       	    match brd.(nr).(nc) with
	      Some p -> if p = player then Some(nr, nc)
		  else findEnd brd player dir nr nc
	    | None -> None
	  end
    	| None -> None in
      fun brd player dir row col ->
    	match (nextSquare dir row col) with
	  Some(r,c) -> begin
       	    match (get brd r c) with
	      Some p ->
 	      	if p = player then None
       	    	else findEnd brd (num_of_piece player) dir r c	    
       	    | None -> None
	  end
    	| None -> None

    let rec legalMove brd player move =
      match move with
    	Pass -> begin (* this is only legal if player has no moves *)
	  try
	    for r = 0 to 7 do
	      for c = 0 to 7 do
 	      	if legalMove brd player (Square(r, c)) then
 	      	  raise Exit
	      done
	    done; true
	  with Exit -> false
    	end
      | Square (r, c) ->
	  if (get brd r c) <> None then
 	    false
	  else
	    List.exists (function dir -> (sandwich brd player dir r c) <> None)
	      (searchDirs r c)

    let makeMove brd player move =
      let playr = num_of_piece player in
      let rec flipPieces b dir r1 c1 r2 c2 =
	if (r1, c1) = (r2, c2) then 0
	else begin
	  b.(r1).(c1) <- playr;
	  let Some(r,c) = nextSquare dir r1 c1 in
	  1 + (flipPieces b dir r c r2 c2)
	end in
      match move with
	Pass -> (copy brd, 0)
      |	Square(row, col) ->
	  let newbrd = copy brd in
	  let flipped = List.fold_left
	      (fun count dir ->
	      	match (sandwich newbrd player dir row col) with
		  None -> count
	      	| Some(r, c) -> count + (flipPieces newbrd dir row col r c))
	      0 (searchDirs row col) in
	  (newbrd, flipped)
	    
    let iter f brd = Array.iter
	(Array.iter (function n -> f (piece_of_num n))) brd
    let iteri f brd =
      for row = 0 to 7 do
	for col = 0 to 7 do
	  f row col (piece_of_num brd.(row).(col))
	done
      done
  end

    
module SquareSet = Set.Make(
  struct
    type t = (int * int)
    let compare (r1, c1) (r2, c2) = ((r1 lsl 3) + c1) - ((r2 lsl 3) + c2)
  end)
    
	
type t = { board: ABoard.t;
 	   turn: piece;		(* whose turn is it next*)
 	   lastMove: move;
 	   bPieces: int;
 	   wPieces: int;
 	   bwDiff: int;
	   freeSquares: SquareSet.t
	 }

let get brd r c = ABoard.get brd.board r c
let set brd r c p = ABoard.set brd.board r c p

let iter f brd = ABoard.iter f brd.board
let iteri f brd = ABoard.iteri f brd.board

let initial () =
  let board = ABoard.create ()
  and free = List.fold_left (fun set squ -> SquareSet.add squ set)
      SquareSet.empty
      [2,2; 2,3; 2,4; 2,5; 3,2; 3,5; 4,2; 4,5; 5,2; 5,3; 5,4; 5,5] in
  let brd = { board=board; bPieces=30; wPieces=30; bwDiff=0; turn=Black;
 	      lastMove=Pass; freeSquares=free } in
  set brd 3 3 (Some White);
  set brd 4 4 (Some White);
  set brd 3 4 (Some Black);
  set brd 4 3 (Some Black);
  brd


let of_array abrd turn last =
  let board = ABoard.create ()
  and free = ref SquareSet.empty
  and black = ref 0
  and white = ref 0 in
  for row = 0 to 7 do
    for col = 0 to 7 do
      ABoard.set board row col abrd.(row).(col);
      match abrd.(row).(col) with
	Some Black -> incr black
      |	Some White -> incr white
      |	None ->
	  List.exists
	    (function (r, c) ->
	      match (ABoard.get board r c) with
		None -> false
	      | Some _ -> free := SquareSet.add (row, col) !free; true)
	    (adjSquares row col); ()
    done
  done;
  let total = !black + !white in
  { board=board;
    bPieces= 32 - (if turn = Black then total / 2 else total / 2 + 1);
    wPieces= 32 - total / 2;
    bwDiff = !black - !white;
    turn = turn;
    lastMove = last;
    freeSquares = !free }
  
let copy brd =
  { board = ABoard.copy brd.board;
    bPieces = brd.bPieces;
    wPieces = brd.wPieces;
    bwDiff = brd.bwDiff;
    turn = brd.turn;
    lastMove = brd.lastMove;
    freeSquares = brd.freeSquares
  }
    
let blackPieces brd = brd.bPieces
let whitePieces brd = brd.wPieces
let bwDiff brd = brd.bwDiff
let turn brd = brd.turn
let lastMove brd = brd.lastMove


let opponent = function
    Black -> White
  | White -> Black

let sandwich brd player dir r c = ABoard.sandwich brd.board player dir r c
let legalMove brd player move = ABoard.legalMove brd.board player move

let canMove brd player = not (legalMove brd player Pass)

let legalMoves brd player =
  let moves =
    SquareSet.fold
      (fun (row, col) mv ->
      	if legalMove brd player (Square(row, col)) then
	  Square(row, col) :: mv
	else mv)
      brd.freeSquares [] in
  if moves = [] then [Pass]
  else moves



let makeMove brd player move =
  match move with
    Pass ->
      { board = ABoard.copy brd.board;
      	turn = opponent player;
      	lastMove = move;
      	bPieces = brd.bPieces;
      	wPieces = brd.wPieces;
      	bwDiff = brd.bwDiff;
      	freeSquares = brd.freeSquares }
  | Square(row, col) ->
      let (nb, count) = ABoard.makeMove brd.board player move in
      { board=nb;
      	turn=opponent player;
      	lastMove=move;
      	bPieces= if player = Black then brd.bPieces - 1 else brd.bPieces;
      	wPieces= if player = White then brd.wPieces - 1 else brd.wPieces;
      	bwDiff= if player = Black then
	  brd.bwDiff + count * 2 + 1 else brd.bwDiff - count * 2 - 1;
      	freeSquares = (List.fold_left
			 (fun set (r, c) ->
			   if (ABoard.get nb r c) = None then
			     SquareSet.add (r, c) set
			   else set)
			 (SquareSet.remove (row, col) brd.freeSquares)
			 (adjSquares row col))
    }

let string_of_piece = function 
    Black -> "BLACK"
  | White -> "WHITE"

let string_of_move = function
    Pass -> "PA"
  | Square(r, c) ->
      [|"A";"B";"C";"D";"E";"F";"G";"H"|].(c) ^ string_of_int(r+1)
  
let printBoard =
  let squareSym = function
    | None -> "-"
    | Some p -> if p = Black then "*" else "O" in
  function brd ->
    let player = turn brd in
    let op = opponent player in
    Printf.printf "%s   (%s)                 A B C D E F G H     %s   (%s)\n"
      (string_of_piece player) (squareSym (Some player))
      (string_of_piece op) (squareSym (Some op));
    for row = 0 to 7 do
      Printf.printf "%26s%d " " " (row + 1);
      for col = 0 to 7 do
 	Printf.printf "%s " (squareSym (get brd row col))
      done;
      Printf.printf "%d\n" (row + 1)
    done;
    Printf.printf "%48sLast move: %s\n" " " (string_of_move (lastMove brd));;
 	 


let printBoardTop =
  let squareSym = function
      None -> "-"
    | Some Black -> "*"
    | Some White -> "O"
  and spc x = String.make x ' ' in
  function brd ->
    let player = turn brd in
    let op = opponent player in
    Format.open_box 0;
    Format.print_string (Printf.sprintf "%-20s"
 			   (Printf.sprintf "%s  (%s)" (string_of_piece player)
 			      (squareSym (Some player))));
    Format.print_string "  A B C D E F G H  ";
    Format.print_string (spc 4);
    Format.print_string (Printf.sprintf "%s  (%s)" (string_of_piece op)
 			   (squareSym (Some op)));
    Format.force_newline();
    for row = 0 to 7 do
      Format.print_string (Printf.sprintf "%20s%d " " " (row + 1));
      for col = 0 to 7 do
 	Format.print_string (Printf.sprintf "%s "
 			       (squareSym (get brd row col)))
      done;
      Format.print_int (row + 1);
      Format.force_newline();
    done;
    Format.print_string ((spc 20) ^ "  A B C D E F G H  " ^ (spc 4));
    Format.print_string (Printf.sprintf "Last move: %s"
 			   (string_of_move (lastMove brd)));
    Format.force_newline();
    Format.close_box()
