(* various feature detection functions *)

(* returns two floats, black's value and white's value *)
type featureFun = Board.t -> (int * int);;

let stablePieces brd =
  let sboard = Array.create_matrix 8 8 None 
  and black = ref 0
  and white = ref 0 in
  (* corner 0, 0 *)
  let lim = ref 7
  and owner = Board.get brd 0 in
  if owner <> None then
    try for row = 0 to 7 do
      if (Board.get brd (row * 8)) = owner then
	try for col = 0 to !lim do
	  if (Board.get brd (row*8 + col)) = owner then
	    sboard.(row).(col) <- owner
	  else begin
	    lim := col - 1;
	    raise Exit
	  end
	done with Exit -> ()
      else
	raise Exit
    done with Exit -> ()
  else ();
  (* corner 0, 7 *)
  let lim = ref 0
  and owner = Board.get brd 7 in
  if owner <> None then
    try for row = 0 to 7 do
      if (Board.get brd (row*8 + 7)) = owner then
	try for col = 7 downto !lim do
	  if (Board.get brd (row*8 + col)) = owner then
	    sboard.(row).(col) <- owner
	  else begin
	    lim := col + 1;
	    raise Exit
	  end
	done with Exit -> ()
      else
	raise Exit
    done with Exit -> ()
  else ();
  (* corner 7, 0 *)
  let lim = ref 7
  and owner = Board.get brd 56 in
  if owner <> None then
    try for row = 7 downto 0 do
      if (Board.get brd (row*8)) = owner then
	try for col = 0 to !lim do
	  if (Board.get brd (row*8 + col)) = owner then
	    sboard.(row).(col) <- owner
	  else begin
	    lim := col - 1;
	    raise Exit
	  end
	done with Exit -> ()
      else
	raise Exit
    done with Exit -> ()
  else ();
  (* corner 7, 7 *)
  let lim = ref 0
  and owner = Board.get brd 63 in
  if owner <> None then
    try for row = 7 downto 0 do
      if (Board.get brd (row*8 + 7)) = owner then
	try for col = 7 downto !lim do
	  if (Board.get brd (row*8 + col)) = owner then
	    sboard.(row).(col) <- owner
	  else begin
	    lim := col + 1;
	    raise Exit
	  end
	done with Exit -> ()
      else
	raise Exit
    done with Exit -> ()
  else ();  
  Array.iter
    (Array.iter (function
	Some Board.Black -> incr black
      | Some Board.White -> incr white
      |	None -> ()))
    sboard;
  (!black, !white);;
	  
let cornerPieces =
  let corners = [0; 7; 56; 63] in
  function brd ->
    List.fold_left
      (fun (black, white) square ->
	match (Board.get brd square) with
	  Some Board.Black -> (black + 1, white)
	| Some Board.White -> (black, white + 1)
	| None -> (black, white))
      (0, 0) corners
   
let xPieces =
  let xes = [9; 14; 49; 54]
  and corners = [0; 7; 56; 63] in
  function brd ->
    List.fold_left2
      (fun (black, white) sq1 sq2 ->
      	match (Board.get brd sq1) with
	  Some Board.Black ->
	    if (Board.get brd sq2) = None then (succ black, white)
	    else (black, white)
	| Some Board.White ->
	    if (Board.get brd sq2) = None then (black, succ white)
	    else (black, white)
	| None -> (black, white))
      (0, 0) xes corners      
  
let pieceCount brd = (Board.blackPieces brd, Board.whitePieces brd)
    
let moveCount brd =
  let black = Board.legalMoves brd Board.Black
  and white = Board.legalMoves brd Board.White in
  ((if black = [Board.Pass] then 0 else List.length black),
   (if white = [Board.Pass] then 0 else List.length white));;
  
let frontierPieces brd =
  XArray.fold_lefti
    (fun i (black, white) ind ->
      let square = Movetable.legalMovesTable.(i).(ind) in
      (black + square.Movetable.bfront,
       white + square.Movetable.wfront))
    (0, 0) (Board.indices brd)

(* "s s+ a+ b+ a+ s* | s* a+ b+ a+ s s+" is the regular expression for an
   edge trap.  Below is the corresponding DFA.
   Q0 = a Q1 | s Q2
   Q1 = a Q1 | b Q3
   Q2 = a Q1 | s Q4
   Q3 = a Q5 | b Q3
   Q4 = a Q6 | s Q4
   Q5 = a Q5 | s Q7
   Q6 = a Q6 | b Q8
   Q7 = s Q9       
   Q8 = a Q10 | b Q8
   Q9 = 1 | s Q9   
   Q10 = 1 | a Q10 | s Q9
   *)  
let edgeTraps brd =
  let black = ref 0
  and white = ref 0 in
   (* counts the traps along a single edge, starting at (row, col) *)
  let rec countTraps dir square state owner =
    match square with
      None -> begin
   	match state with
	  9 | 10 -> if owner = Board.Black then incr black else incr white
   	| _ -> ()
      end
    | Some sq -> begin
	let next = Board.nextSquare dir sq in
    	match (Board.get brd sq) with
	  Some p -> begin
	    match state with
	      0 -> countTraps dir next 1 p
	    | 1 -> if p = owner then countTraps dir next 1 owner
	           else countTraps dir next 3 owner
	    | 2 -> countTraps dir next 1 p
	    | 3 -> if p = owner then countTraps dir next 5 owner
		   else countTraps dir next 3 owner
	    | 4 -> countTraps dir next 6 p
	    | 5 -> if p = owner then countTraps dir next 5 owner
		   else ()
	    | 6 -> if p = owner then countTraps dir next 6 owner
		   else countTraps dir next 8 owner
	    | 7 -> ()
	    | 8 -> if p = owner then countTraps dir next 10 owner
		   else countTraps dir next 8 owner
	    | 9 -> ()
	    | 10 -> if p = owner then countTraps dir next 10 owner
		    else ()
	    | _ -> raise (Invalid_argument "Invalid state in edgeTraps")
	  end
	| None -> begin
	    match state with
	      0 -> countTraps dir next 2 owner
	    | 2 | 4 -> countTraps dir next 4 owner
	    | 5 -> countTraps dir next 7 owner
	    | 7 | 9 | 10 -> countTraps dir next 9 owner
	    | 1 | 3 | 6 | 8 -> ()
	    | _ -> raise (Invalid_argument "Invalid state in edgeTraps")
	end
    end in
  countTraps 2 (Some 0) 0 Board.Black;
  countTraps 2 (Some 56) 0 Board.Black;
  countTraps 4 (Some 0) 0 Board.Black;
  countTraps 4 (Some 7) 0 Board.Black;
  (!black, !white);;

let testFeatures brd =
  let testFeature (f, name) =
    let (b, w) = f brd in
    Printf.printf "%s: (%d, %d)\n" name b w in
  List.iter testFeature [ stablePieces, "stablePieces";
			  cornerPieces, "cornerPieces";
			  xPieces, "xPieces";
			  pieceCount, "pieceCount";
			  moveCount, "moveCount";
			  frontierPieces, "frontierPieces";
			  edgeTraps, "edgeTraps"];;
  
