{-# LANGUAGE BangPatterns #-}
-- | Board State 
module Board
  ( -- * Types
    Board (..)
  , PieceMap
  , Color (..)
  , Kind (..)
  , Piece
  , Position (..)
  , Move (..)
  , Game (..)
  , initBoard
  , initGame
  , color, kind, height   -- attributes of pieces 
  , nthTurn, nthMove
  , invert
  , inactive
  , countPieces
  , endGame
  , winner
  , swapBoard
  , captureMoves    
  , stackingMoves
  , nextMoves
  , countStacks
  , sumHeights
  , maxHeights
  , sixLines
  , emptyBoard
  , startingBoard
  , randomBoard
  , randomBoardIO
  , showMove
  , applyMove
  , applyMoveSkip
  , applyMoveSkip'
  , positions
  , zoneOfControl
  ) where

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


-- | player colors
data Color = White | Black 
           deriving (Eq,Show,Enum,Read)

-- | the inverse color
invert :: Color -> Color
invert White = Black
invert Black = White

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

-- | A piece stack: color, kind and height (starting at 1).
type Piece = (Color, Kind, Int)

color :: Piece -> Color
color (c, _, _) = c

kind :: Piece -> Kind
kind (_, k, _)= k

height :: Piece -> Int
height (_, _, h)= h


-- | Algebraic board positions.  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 (Eq, Ord, Enum, Bounded, Show, Read)

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

-- | A mapping from positions to pieces 
type PieceMap = Map Position Piece

-- | The board state
-- | current turn, active player pieces, other player pieces
data Board 
      = Board { active :: !Color,            -- player to move next 
                move :: !Int,                -- total move count
                moves :: [Move],             -- all previous moves (in reverse order)
                pieces :: PieceMap,          -- map positions to pieces
                activeCounts :: [Int],       -- active player counts
                inactiveCounts :: [Int],     -- inactive player counts
                activeHeights :: [Int],
                inactiveHeights :: [Int]
              } deriving (Eq, Show, Read)


-- | initialize a board from a list of piece & positions
initBoard :: [(Position,Piece)] -> Board
initBoard assocs 
  = let ps = Map.fromList assocs
    in Board { active=White, move=0, moves=[], pieces=ps,
               activeCounts=countStacks White ps,
               inactiveCounts=countStacks Black ps,
               activeHeights=sumHeights White ps,
               inactiveHeights=sumHeights Black ps 
             }
       
inactive :: Board -> Color
inactive = invert . active



-- | A move is either a capture, a stacking or a pass
--   "Skip" is a dummy move to alternate players in a turn
data Move = Capture !Position !Position  -- origin and destination positions
          | Stack   !Position !Position  
          | Pass 
          | Skip
            deriving (Eq, Show, Read)


showMove :: Move -> String
showMove (Capture a b) = show a ++ "x" ++ show b
showMove (Stack a b)   = show a ++ "-" ++ show b
showMove Pass          = "pass"
showMove Skip          = "skip"


-- | record to hold the game state
data Game = Game
  { human :: Color       -- human plays white or black?
  , initial :: Board     -- initial board
  , board :: Board       -- current board
  , trail :: [Move]      -- previous moves 
  , remain :: [Move]     -- future moves
  } deriving (Eq,Show) 


-- | initialize a game state
initGame :: Board -> Color -> Game
initGame b c = Game { human=c, initial=b, board=b, trail=[], remain=[] }



-- | Convert number of moves into number of turns 
nthTurn :: Int -> Int
nthTurn 0 = 1
nthTurn m | m>0 = 2 + (m-1)`div`3

nthMove :: Int -> Int
nthMove 0 = 1
nthMove m | m>0 = min (1+(m-1)`mod`3) 2



-- | number of pieces in a board
countPieces :: Board -> Int
countPieces board = Map.size (pieces board) 




-- | swap active player
-- this is used to analyse the board from the opponent's prespective
swapBoard :: Board -> Board
swapBoard b = b { active= invert (active b), 
                  activeCounts = inactiveCounts b, 
                  inactiveCounts = activeCounts b,
                  activeHeights= inactiveHeights b,
                  inactiveHeights = activeHeights b
                }


-- | next moves for the active player
nextMoves :: Board -> [Move]
nextMoves b 
  | tzaars==0 || tzarras==0 || totts==0 = []
  | m == 0                              = captureMoves b
  | otherwise 
    = case (m-1)`mod`3 of
        0 -> captureMoves b     -- first move
        1 -> [Skip]             -- dummy opponent move within a turn
        2 -> stackingMoves b ++ captureMoves b ++ [Pass] -- second moves
        _ -> error "nextMoves: invalid board"
  where -- lostPieces = any (==0) (countStacks (active b) (pieces b)) 
        tzaars:tzarras:totts:_ = activeCounts b
        m = move b



-- | next capture moves for the active player
captureMoves :: Board -> [Move]
captureMoves board = Map.foldrWithKey forPiece [] (pieces board)
  where
    c = active board
    forPiece :: Position -> Piece -> [Move] -> [Move]
    forPiece !p (!c', _, !i) moves 
      | c==c'     = foldl' downLine moves (sixLines p)
      | otherwise = moves
        where
          downLine :: [Move] -> [Position] -> [Move]
          downLine moves []   = moves
          downLine moves (q:ps) 
              = case Map.lookup q (pieces board) of
                  Nothing -> downLine moves ps 
                  Just (c', _, j) | c/=c' && i>=j -> (Capture p q):moves
                  _  -> moves



-- | next stacking moves for the active player
stackingMoves :: Board -> [Move]
stackingMoves board = foldl' forPiece [] (Map.assocs (pieces board))
  where 
    c = active board
    tzaars:tzarras:totts:_ = activeCounts board
    forPiece :: [Move] -> (Position,Piece) -> [Move]
    forPiece moves (p,(c',_,_)) 
      | c==c'     = foldl' downLine moves (sixLines p)
      | otherwise = moves
        where
          downLine :: [Move] -> [Position] -> [Move]
          downLine moves [] = moves
          downLine moves (q:ps) 
              = case Map.lookup q (pieces board) of
                  Nothing  -> downLine moves ps
                  Just (c', _, _) | c/=c' -> moves
                  Just (_, Tzaar, _) | tzaars==1  -> moves
                  Just (_, Tzarra, _) | tzarras==1 -> moves
                  Just (_, Tott, _) | totts==1  -> moves
                  Just (_, _, _)  -> (Stack p q) : moves




-- | count the number of stacks of each type in a half-board
countStacks :: Color -> PieceMap -> [Int]
countStacks c pieces
    = count 0 0 0 (Map.elems pieces)
    where
      count :: Int -> Int -> Int -> [Piece] -> [Int]
      count !x !y !z ((c',Tzaar,_) : ps)  | c==c' = count (1+x) y z ps
      count !x !y !z ((c',Tzarra,_) : ps) | c==c' = count x (1+y) z ps
      count !x !y !z ((c',Tott,_) : ps)   | c==c' = count x y (1+z) ps
      count !x !y !z (_ : ps)                     = count x y z ps      
      count !x !y !z []              = [x,y,z]


-- | sum of heights of stacks for each kind
sumHeights :: Color -> PieceMap -> [Int]
sumHeights c pieces = sum 0 0 0 (Map.elems pieces)
  where sum :: Int -> Int -> Int -> [Piece] -> [Int]
        sum !x !y !z ((c',Tzaar,!h):ps)  | c==c' = sum (x+h) y z ps
        sum !x !y !z ((c',Tzarra,!h):ps) | c==c' = sum x (y+h) z ps
        sum !x !y !z ((c',Tott,!h):ps)   | c==c' = sum x y (z+h) ps
        sum !x !y !z (_ : ps)        = sum x y z ps
        sum !x !y !z []              = [x,y,z]

-- | maximum height for each kind
maxHeights :: Color -> PieceMap -> [Int]        
maxHeights c pieces = maxh 0 0 0 (Map.elems pieces)
  where maxh :: Int -> Int -> Int -> [Piece] -> [Int]
        maxh !x !y !z ((c',Tzaar,!h):ps) | c==c' && h>x = maxh h y z ps
        maxh !x !y !z ((c',Tzarra,!h):ps) | c==c' && h>y= maxh x h z ps        
        maxh !x !y !z ((c',Tott,!h):ps) | c==c' && h>z  = maxh x y h ps
        maxh !x !y !z (_ : ps) = maxh x y z ps
        maxh !x !y !z []       = [x,y,z]


-- | The next board state after a move.  
-- | Assumes the move is valid.
applyMove :: Move -> Board -> Board
applyMove m@(Capture x y) b 
  = b {active=invert (active b), 
       move=1+move b, 
       moves = m:moves b,
       pieces= pieces', 
       activeCounts = counts',         -- swap counts
       inactiveCounts= activeCounts b,
       activeHeights = heights',       -- swap heights
       inactiveHeights = activeHeights b
      }
    where
      pX  = pieces b!x      
      (_, kindY, sizeY) = pieces b!y
      pieces' = Map.insert y pX (Map.delete x (pieces b))
      counts' = increment kindY (-1)     (inactiveCounts b)
      heights'= increment kindY (-sizeY) (inactiveHeights b)


applyMove m@(Stack x y) b 
  = b {active=invert (active b), 
       move=1+move b, 
       moves= m:moves b,
       pieces=pieces',
       activeCounts = inactiveCounts b,
       inactiveCounts = counts',
       activeHeights = inactiveHeights b,
       inactiveHeights= heights'
         }
    where
      (colorX, kindX, sizeX) = pieces b!x
      (_,      kindY, sizeY) = pieces b!y
      pieces' = Map.insert y (colorX,kindX,sizeX+sizeY) (Map.delete x (pieces b))
      counts' = increment kindY (-1) (activeCounts b)
      heights' | kindX==kindY = activeHeights b
               | otherwise = increment kindY (-sizeY) $ 
                             increment kindX sizeY (activeHeights b)
          
          
-- Pass & Skip have the same effect
applyMove m b = b {active= invert (active b), 
                   move=1+move b, 
                   moves= m:moves b,
                   activeCounts= inactiveCounts b,
                   inactiveCounts= activeCounts b,
                   activeHeights= inactiveHeights b,
                   inactiveHeights= activeHeights b
                  }

-- | modify a counter
increment :: Kind -> Int -> [Int] -> [Int]
increment Tzaar  i (tzaars:tzarras:totts:_) = (tzaars+i) : tzarras : totts : []
increment Tzarra i (tzaars:tzarras:totts:_) = tzaars : (tzarras+i) : totts : []
increment Tott   i (tzaars:tzarras:totts:_) = tzaars : tzarras : (totts+i) : []


-- | apply one move and subsequent skip move
applyMoveSkip :: Move -> Board -> Board 
applyMoveSkip m b 
  = case nextMoves b' of
    (Skip:_) -> applyMove Skip b'  -- compulsory move
    _      -> b'
    where b' = applyMove m b


-- | apply one move checking for validity
applyMoveSkip' :: Board -> Move -> Maybe Board
applyMoveSkip' b m
  | m`elem`ms = 
    case nextMoves b' of
      (Skip:_) -> Just (applyMove Skip b')
      _ -> Just b'
  | otherwise = Nothing
  where ms = nextMoves b
        b' = applyMove m b

                          
endGame :: Board -> Bool
endGame = null . nextMoves 

winner :: Board -> Color
winner = invert . active


-- | 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 map should be memoied lazily 
sixLines_memo :: Map Position [[Position]]  -- 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 = sixLines_memo!p




-- | An empty board
emptyBoard :: Board
emptyBoard = initBoard [] 

-- | The default (non-randomized, non-tournament) starting position.
startingBoard :: Board
startingBoard   = initBoard (whites ++ blacks) 
  where
  whites = [(p, (White,Tzaar,1)) | p<-wTzaars] ++ 
           [(p, (White,Tzarra,1)) | p<-wTzarras] ++ 
           [(p, (White,Tott,1)) | p<-wTotts]
  blacks = [(p, (Black,Tzaar,1)) | p<-bTzaars] ++ 
           [(p, (Black,Tzarra,1)) | p<-bTzarras] ++ 
           [(p, (Black,Tott,1)) | p<-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 = (b, rnd')
    where b = initBoard (whites++blacks) 
          ws = replicate 6 (White,Tzaar,1) ++
               replicate 9 (White,Tzarra,1) ++
               replicate 15 (White,Tott,1)
          bs = replicate 6 (Black,Tzaar,1) ++
               replicate 9 (Black,Tzarra,1) ++
               replicate 15 (Black,Tott,1)
          (positions',rnd') = shuffle rnd positions
          whites = zip (take 30 positions') ws
          blacks = zip (drop 30 positions') bs

randomBoardIO :: IO Board
randomBoardIO = do rnd <- getStdGen
                   let (b, rnd') = randomBoard rnd
                   setStdGen rnd'
                   return b


-- 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)




-- Estimate the zone of control of a player
-- i.e., the opponents' pieces that can be captured in two moves
zoneOfControl ::  Color -> PieceMap -> PieceMap
zoneOfControl c pieces = Map.filterWithKey forPiece1 pieces
    where
      -- player's pieces that make at least one capture
      movable = Map.filterWithKey forPiece2 pieces

      forPiece1, forPiece2 :: Position -> Piece -> Bool
      forPiece1 p (c', _, i) = c'/=c && or (map (downLine0 i) $ sixLines p)
      forPiece2 p (c',_, h) = c'==c && or (map (downLine2 h) $ sixLines p)

      downLine0, downLine1, downLine2 :: Int -> [Position] -> Bool

      downLine0 i [] = False
      downLine0 i (p:ps) 
          = case Map.lookup p pieces of
              Nothing -> downLine0 i ps
              Just (c', _, h) | c'==c -> 
                  h>=i || (p`Map.member`movable && downLine1 i ps)
              Just (c', _, j) | c'/=c -> 
                  or $ map (downLine1 (max i j)) $ sixLines p

      downLine1 i [] = False
      downLine1 i (p:ps) 
          = case Map.lookup p pieces of
              Nothing -> downLine1 i ps
              Just (c', _, h) -> c'==c && h>=i

      downLine2 h [] = False
      downLine2 h (p:ps) 
          = case Map.lookup p pieces of
              Nothing -> downLine2 h ps
              Just (c', _, i) -> c'/=c && h>=i