-- | Board State and AI
module Board
  (
  -- * Types
    Board
  , HalfBoard
  , BoardTree
  , GameTree(..)
  , Type (..)
  , Piece
  , Position (..)
  , Move
  , Turn
  -- , AtPosition
  , Strategy
  , AI (..)
  -- * Utilities
  , boardTree
  , swapBoard
  , swapBoardTree
  , nextCaptureMoves
  , nextStackingMoves
  , nextTurns
  , connectedPositions
  , threeLines
  , sixLines
  , atPosition
  , startingBoard
  , randomBoard
  , showTurn
  , showMove
  , applyMove
  , applyTurn
  ) where

import Data.List
import Data.Map (Map)
import qualified Data.Map as Map
import System.Random
import Control.Monad(mplus)

-- | The board state is a pair of two "half-boards" (one per player)
type Board = (HalfBoard, HalfBoard)

-- | A Half-board maps locations to pieces 
type HalfBoard = Map Position Piece

-- | A game tree with nodes s and moves m
data GameTree s m = GameTree s [(m, GameTree s m)] deriving Show

-- | A game tree of boards labeled with a boolean 
--   the label is True if your turn, False if opponent.
type BoardTree = GameTree (Bool,Board) Turn

-- | The three types of pieces
-- | Each player starts with 6 Tzaars, 9 Tzarras, and 15 Totts.
data Type = Tzaar | Tzarra | Tott deriving (Show, Eq)

-- | the type of a piece, and the level of the stack (starting with 1).
type Piece = (Type, Int)

-- | Board position.  Letters left to right, numbers bottom to top.
--   Column E has the hole in the middle.
data Position
  = A1 | A2 | A3 | A4 | A5
  | B1 | B2 | B3 | B4 | B5 | B6
  | C1 | C2 | C3 | C4 | C5 | C6 | C7
  | D1 | D2 | D3 | D4 | D5 | D6 | D7 | D8
  | E1 | E2 | E3 | E4 | E5 | E6 | E7 | E8
  | F1 | F2 | F3 | F4 | F5 | F6 | F7 | F8
  | G1 | G2 | G3 | G4 | G5 | G6 | G7
  | H1 | H2 | H3 | H4 | H5 | H6
  | I1 | I2 | I3 | I4 | I5
  deriving (Show, Eq, Ord, Enum, Bounded)

-- | A move is one position to another, for either capturing or stacking.
type Move = (Position, Position)

-- | A complete turn is move, followed by an optional move.
type Turn = (Move, Maybe Move)

-- | An AI strategy calculates the next turn from a board tree.
type Strategy = BoardTree -> StdGen -> (Turn, StdGen)

-- | An AI player.
data AI = AI
  { name        :: String   -- ^ Name of AI.
  , description :: String   -- ^ Brief description of AI.
  , strategy    :: Strategy -- ^ The strategy.
  }

-- | The state of a single board position; Right true if you, Left if opponent.
-- type AtPosition = Either Piece Piece

-- | List of all positions (for enumeration purposes)
positions :: [Position]
positions = [minBound .. maxBound]

showTurn :: Turn -> String
showTurn (a, Nothing) = showMove a
showTurn (a, Just b ) = showMove a ++ "    " ++ showMove b

showMove :: Move -> String
showMove (a, b) = show a ++ " -> " ++ show b



-- | Possible next turns.
nextTurns :: Board -> [Turn]
nextTurns board@(you, _)
  | lostOneOfThree = []
  | otherwise      = captureCapture ++ captureStack ++ captureNothing
  where
  a = nextCaptureMoves board
  b = map (applyMove board) a
  c = map nextCaptureMoves  b
  d = map nextStackingMoves b
  captureCapture = [ (a, Just b) | (a, x) <- zip a c, b <- x ]
  captureStack   = [ (a, Just b) | (a, x) <- zip a d, b <- x ]
  captureNothing = zip a $ repeat Nothing
  lostOneOfThree = length (nub [t | (t, _)<-Map.elems you]) /= 3


nextCaptureMoves :: Board -> [Move]
nextCaptureMoves board@(you, _) = concatMap forPiece (Map.assocs you)
  where
  forPiece :: (Position,Piece) -> [Move]
  forPiece (p, (_, i)) = concatMap downLine $ sixLines p
    where
    downLine :: [Position] -> [Move]
    downLine [] = []
    downLine (a:b) = case atPosition board a of
      Nothing -> downLine b
      Just (True, _) -> []
      Just (False, (_, j)) -> [(p, a) | i>=j]

nextStackingMoves :: Board -> [Move]
nextStackingMoves board@(you, _) = concatMap forPiece (Map.keys you)
  where
  forPiece :: Position -> [Move]
  forPiece p = concatMap downLine $ sixLines p
    where
    downLine :: [Position] -> [Move]
    downLine [] = []
    downLine (a:b) = case atPosition board a of
      Nothing   -> downLine b
      Just (False, _) -> []
      Just (True, (Tzaar,_)) | oneTzaarRemaining  -> []
      Just (True, (Tzarra,_)) | oneTzarraRemaining -> []
      Just (True, (Tott, _)) | oneTottRemaining   -> []
      Just (True, _) -> [(p, a)]
  oneTzaarRemaining  = 1 == Map.size (Map.filter (\(t,_)->t==Tzaar) you)
  oneTzarraRemaining = 1 == Map.size (Map.filter (\(t,_)->t==Tzarra) you)
  oneTottRemaining   = 1 == Map.size (Map.filter (\(t,_)->t==Tott) you)





-- Creates a board tree for you and opponent.  Assumes you have the next turn.
boardTree :: Board -> BoardTree
boardTree board = mkTree True board
  where
  mkTree :: Bool -> Board -> BoardTree
  mkTree you b
      = GameTree (you,if you then b else swapBoard b) 
            [ (t, mkTree (not you) $ swapBoard $ applyTurn b t) 
                  | t<-nextTurns b]




-- | Swaps board positions, i.e. white to black, black to white.
swapBoard :: Board -> Board
swapBoard (a, b) = (b, a)

-- | Swaps board trees, i.e. white to black, black to white.
swapBoardTree :: BoardTree -> BoardTree
swapBoardTree (GameTree (you,board) branches) = GameTree (not you,swapBoard board) [ (t, swapBoardTree bt) | (t, bt) <- branches ]


-- Querying the state of a board position.
atPosition :: Board -> Position -> Maybe (Bool,Piece)
atPosition (you,opp) pos 
    = do { piece<-Map.lookup pos you
         ; return (True,piece) 
         } `mplus`
      do { piece<-Map.lookup pos opp
         ; return (False,piece)
         }


-- | All the lines that form connected positions on the board.
connectedPositions :: [[Position]]
connectedPositions =
  [ [A1, A2, A3, A4, A5]
  , [B1, B2, B3, B4, B5, B6]
  , [C1, C2, C3, C4, C5, C6, C7]
  , [D1, D2, D3, D4, D5, D6, D7, D8]
  , [E1, E2, E3, E4]
  , [E5, E6, E7, E8]
  , [F1, F2, F3, F4, F5, F6, F7, F8]
  , [G1, G2, G3, G4, G5, G6, G7]
  , [H1, H2, H3, H4, H5, H6]
  , [I1, I2, I3, I4, I5]

  , [A1, B1, C1, D1, E1]
  , [A2, B2, C2, D2, E2, F1]
  , [A3, B3, C3, D3, E3, F2, G1]
  , [A4, B4, C4, D4, E4, F3, G2, H1]
  , [A5, B5, C5, D5]
  , [F4, G3, H2, I1]
  ,     [B6, C6, D6, E5, F5, G4, H3, I2]
  ,         [C7, D7, E6, F6, G5, H4, I3]
  ,             [D8, E7, F7, G6, H5, I4]
  ,                 [E8, F8, G7, H6, I5]
   
  ,                 [E1, F1, G1, H1, I1]
  ,             [D1, E2, F2, G2, H2, I2]
  ,         [C1, D2, E3, F3, G3, H3, I3]
  ,     [B1, C2, D3, E4, F4, G4, H4, I4]
  , [A1, B2, C3, D4]
  , [F5, G5, H5, I5]
  , [A2, B3, C4, D5, E5, F6, G6, H6]
  , [A3, B4, C5, D6, E6, F7, G7]
  , [A4, B5, C6, D7, E7, F8]
  , [A5, B6, C7, D8, E8]
  ]

-- | The three lines that cross at a single board position.
threeLines :: Position -> [[Position]]
threeLines p = [ line | line <- connectedPositions, elem p line ]


-- | The six lines traveling radially out from a single board position.
-- | optimization: this function is lazily memoied 

sixLines_memo :: Map Position [[Position]]
sixLines_memo = Map.fromList [(p, radials p) | p<-positions]
    where radials p = [r | l<-threeLines p, r<-divide p l, not (null r)]
          divide a b = [reverse x, y]
              where (x, _:y) = span (/= a) b

sixLines :: Position -> [[Position]]
sixLines p = Map.findWithDefault undefined p sixLines_memo





-- | The next board state after a move.  Assumes move is valid.
applyMove :: Board -> Move -> Board
applyMove board@(a, b) (x, y) 
    | fromA     = (Map.insert y piece (Map.delete x a'), b')
    | otherwise = (a', Map.insert y piece (Map.delete x b'))
    where
      Just (whoX, (typeX, sizeX)) = atPosition board x
      Just (whoY, (_    , sizeY)) = atPosition board y
      capture = whoX /= whoY
      fromA = Map.member x a
      piece = (typeX, if capture then sizeX else sizeX + sizeY)
      a' = Map.delete y a
      b' = Map.delete y b

-- | The next board state after a complete turn.  Assumes turn is valid.
applyTurn :: Board -> Turn -> Board
applyTurn board (a, Just b ) = applyMove (applyMove board a) b
applyTurn board (a, Nothing) =            applyMove board a



-- | The default (non-randomized, non-tournament) starting position.
startingBoard :: Board
startingBoard = (Map.fromList whites, Map.fromList blacks)
  where
  f t p = (p, (t, 1))
  whites = map (f Tzaar) wTzaars ++ map (f Tzarra) wTzarras ++ map (f Tott) wTotts
  blacks = map (f Tzaar) bTzaars ++ map (f Tzarra) bTzarras ++ map (f Tott) bTotts
  wTzaars  = [D3, E3, G4, G5, C5, D6]
  wTzarras = [C2, D2, E2, H3, H4, H5, B5, C6, D7]
  wTotts   = [B1, C1, D1, E1, I2, I3, I4, I5, D8, C7, B6, A5, E4, F5, D5]
  bTzaars  = [C3, C4, F3, G3, E6, F6]
  bTzarras = [B2, B3, B4, F2, G2, H2, E7, F7, G6]
  bTotts   = [A1, A2, A3, A4, F1, G1, H1, I1, E8, F8, G7, H6, D4, E5, F4]


-- | A randomized starting position
randomBoard :: StdGen -> (Board, StdGen)
randomBoard rnd 
    = ((Map.fromList whites, Map.fromList blacks), rnd')
    where pieces = replicate 6 (Tzaar,1) ++
                   replicate 9 (Tzarra,1) ++
                   replicate 15 (Tott,1)
          (positions',rnd') = shuffle rnd positions
          whites = zip (take 30 positions') pieces
          blacks = zip (drop 30 positions') pieces


-- an auxilary function to shuffle a list randomly
shuffle :: StdGen  -> [a] -> ([a], StdGen)
shuffle g xs = shuffle' g xs (length xs)
    where
      shuffle' :: RandomGen g => g -> [a] -> Int -> ([a], g)
      shuffle' g xs n 
          | n>0 = let (k, g') = randomR (0,n-1) g
                      (xs',x:xs'') = splitAt k xs
                      (ys,g'') = shuffle' g' (xs' ++ xs'') (n-1)
                  in (x:ys, g'')
          | otherwise = ([],g)