{-# LANGUAGE BangPatterns, RecordWildCards #-}
-- | 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
  , swapPlayer
  , captureMoves
  , stackingMoves
  , nextMoves
  , countStacks
  , sumHeights
  , maxHeights
  , sixLines
  , emptyBoard
  , startingBoard
  , randomBoard
  , randomBoardIO
  , showMove
  , applyMove
  , applyMoveSkip
  , positions
  , zoneOfControl
  ) where

import           Data.List (foldl')
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import           Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import           Data.Hashable

import qualified Data.Vector.Unboxed as Vec
import           Data.Vector.Unboxed(Vector)

import           Data.Array

import           Control.Monad (liftM, mplus)
import           System.Random

-- | 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, Enum, Show, Read)

-- | A piece stack: color, kind and height (starting at 1).
data Piece = Piece !Color !Kind !Int
             deriving (Eq, Show, Read)

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

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

height :: Piece -> Int
height (Piece _ _ 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 (Ix, Eq, Ord, Enum, Bounded, Show, Read)

instance Hashable Position where
  hashWithSalt s p = hashWithSalt s (fromEnum p)


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

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

-- | A TZAAR game board
-- | current turn, active player pieces, other player pieces
data Board
      = Board { active :: !Color,            -- player to move next
                move :: !Int,                -- total move count
                pieces :: PieceMap,          -- map positions to pieces
                activeCounts :: !(Vector Int),       -- active player counts
                inactiveCounts :: !(Vector Int),     -- inactive player counts
                activeHeights :: !(Vector Int),
                inactiveHeights :: !(Vector Int)
              } deriving (Eq, Show, Read)


-- | initialize a board from a list of piece & positions
initBoard :: [(Position,Piece)] -> Board
initBoard assocs
  = let ps = HashMap.fromList assocs
    in Board { active=White,  move=0, 
               pieces=ps,
               activeCounts=Vec.fromList (countStacks White ps),
               inactiveCounts=Vec.fromList (countStacks Black ps),
               activeHeights=Vec.fromList (sumHeights White ps),
               inactiveHeights=Vec.fromList (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 = HashMap.size (pieces board)




-- | swap active player
swapPlayer :: Board -> Board
swapPlayer b = b { active = invert (active b),
                     activeCounts = inactiveCounts b,
                     inactiveCounts = activeCounts b,
                     activeHeights = inactiveHeights b,
                     inactiveHeights = activeHeights b
                   }


-- | all available moves for the active player
nextMoves :: Board -> [Move]
nextMoves b
  | Vec.any (==0) (activeCounts b) = []
  | 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 m = move b



-- | next capture moves for the active player
captureMoves :: Board -> [Move]
captureMoves Board{..} = HashMap.foldrWithKey forPiece [] pieces
  where
    forPiece :: Position -> Piece -> [Move] -> [Move]
    forPiece !p (Piece c _ i) moves
      | c==active = foldl' downLine moves (sixLines p)
      | otherwise = moves
        where
          downLine :: [Move] -> [Position] -> [Move]
          downLine moves []   = moves
          downLine moves (q:rest)
              = case {-# SCC "piece-lookup" #-} HashMap.lookup q pieces of
                  Nothing -> downLine moves rest
                  Just (Piece c _ j) | c/=active && i>=j -> (Capture p q):moves
                  _  -> moves



-- | next stacking moves for the active player
stackingMoves :: Board -> [Move]
stackingMoves board = HashMap.foldlWithKey' forPiece [] (pieces board)
  where
    c = active board
    tzaars = activeCounts board Vec.! fromEnum Tzaar
    tzarras= activeCounts board Vec.! fromEnum Tzarra
    totts  = activeCounts board Vec.! fromEnum Tott
    forPiece :: [Move] -> Position -> Piece -> [Move]
    forPiece moves p (Piece c' _ _)
      | c==c'     = foldl' downLine moves (sixLines p)
      | otherwise = moves
        where
          downLine :: [Move] -> [Position] -> [Move]
          downLine moves [] = moves
          downLine moves (q:ps)
              = case {-# SCC "piece-lookup" #-} HashMap.lookup q (pieces board) of
                  Nothing  -> downLine moves ps
                  Just (Piece c' _ _) | c/=c' -> moves
                  Just (Piece _ Tzaar _) | tzaars==1  -> moves
                  Just (Piece _ Tzarra _) | tzarras==1 -> moves
                  Just (Piece _ 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 (HashMap.elems pieces)
    where
      count :: Int -> Int -> Int -> [Piece] -> [Int]
      count !x !y !z ((Piece c' Tzaar _) : ps)  | c==c' = count (1+x) y z ps
      count !x !y !z ((Piece c' Tzarra _) : ps) | c==c' = count x (1+y) z ps
      count !x !y !z ((Piece 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 (HashMap.elems pieces)
  where sum :: Int -> Int -> Int -> [Piece] -> [Int]
        sum !x !y !z ((Piece c' Tzaar !h):ps)  | c==c' = sum (x+h) y z ps
        sum !x !y !z ((Piece c' Tzarra !h):ps) | c==c' = sum x (y+h) z ps
        sum !x !y !z ((Piece 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 (HashMap.elems pieces)
  where maxh :: Int -> Int -> Int -> [Piece] -> [Int]
        maxh !x !y !z ((Piece c' Tzaar !h):ps) | c==c' && h>x = maxh h y z ps
        maxh !x !y !z ((Piece c' Tzarra !h):ps) | c==c' && h>y= maxh x h z ps
        maxh !x !y !z ((Piece 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,
       pieces= pieces',
       activeCounts = counts',         -- swap counts
       inactiveCounts= activeCounts b,
       activeHeights = heights',       -- swap heights
       inactiveHeights = activeHeights b
      }
    where
      pX  = pieces b HashMap.! x
      (Piece _ kindY sizeY) = pieces b HashMap.! y
      pieces' = {-# SCC "piece-insert" #-} HashMap.insert y pX (HashMap.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,
       pieces=pieces',
       activeCounts = inactiveCounts b,
       inactiveCounts = counts',
       activeHeights = inactiveHeights b,
       inactiveHeights= heights'
         }
    where
      (Piece colorX kindX sizeX) = pieces b HashMap.! x
      (Piece _     kindY sizeY) = pieces b HashMap.! y
      pieces' = {-# SCC "piece-ins-del" #-} HashMap.insert y (Piece colorX kindX (sizeX+sizeY)) (HashMap.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,
                   activeCounts= inactiveCounts b,
                   inactiveCounts= activeCounts b,
                   activeHeights= inactiveHeights b,
                   inactiveHeights= activeHeights b
                  }

-- | modify a counter
increment :: Kind -> Int -> Vector Int -> Vector Int
increment !k !n v = v Vec.// [(i, n + v Vec.! i)]
  where !i = fromEnum k


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



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 six lines traveling radially out from a single board position.
sixLines :: Position -> [[Position]]
sixLines p = sixLinesArray ! p

-- | global memoied array  
sixLinesArray :: Array Position [[Position]]
sixLinesArray = array (minBound,maxBound) [(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

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

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

-- | The default (non-randomized, non-tournament) starting position.
startingBoard :: Board
startingBoard   = initBoard (whites ++ blacks)
  where
  whites = [(p, (Piece White Tzaar 1)) | p<-wTzaars] ++
           [(p, (Piece White Tzarra 1)) | p<-wTzarras] ++
           [(p, (Piece White Tott 1)) | p<-wTotts]
  blacks = [(p, (Piece Black Tzaar 1)) | p<-bTzaars] ++
           [(p, (Piece Black Tzarra 1)) | p<-bTzarras] ++
           [(p, (Piece 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 (Piece White Tzaar 1) ++
               replicate 9 (Piece White Tzarra 1) ++
               replicate 15 (Piece White Tott 1)
          bs = replicate 6 (Piece Black Tzaar 1) ++
               replicate 9 (Piece Black Tzarra 1) ++
               replicate 15 (Piece 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 = HashMap.filterWithKey forPiece1 pieces
    where
      -- player's pieces that make at least one capture
      movable = HashMap.filterWithKey forPiece2 pieces

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

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

      downLine0 !i [] = False
      downLine0 !i (p:ps)
          = case HashMap.lookup p pieces of
              Nothing -> downLine0 i ps
              Just (Piece c' _ h) -> if c==c' then
                    h>=i || (p`HashMap.member`movable && downLine1 i ps)
                  else
                    let !h' = max i h
                    in any (downLine1 h') (sixLines p)


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

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