module AI.Minimax(greedy, ply2,ply3,ply4, dyn1, dyn2) where

import Data.List (sort, sortBy, maximumBy, minimumBy, nub, nubBy)
import qualified Data.Map as Map
import Data.Map (Map)
import AI.Utils
import Board
import Debug.Trace


greedy :: AI
greedy = AI { name = "greedy"
            , description = "Maximize the static evaluation function"
            , strategy = (ifPieces (==60) 
                          (firstTurn greedyStrategy)
                          (ifPieces (>52)
                           (onlyCaptureStack greedyStrategy)
                           (narrowDoubleCaptures greedyStrategy)
                          )
                         )
            }


ply2 :: AI
ply2 = AI { name = "ply2"
          , description = "Minimax with depth 2"
          , strategy = (ifPieces (==60) 
                        (firstTurn $ minimaxStrategy 2 3)
                        (narrowDoubleCaptures $ minimaxStrategy 2 3)
                       )
          }

ply3 :: AI
ply3 = AI { name = "ply3"
          , description = "Minimax with depth 3"
          , strategy = (ifPieces (==60) 
                        (firstTurn $ minimaxStrategy 3 3)
                        (narrowDoubleCaptures $ minimaxStrategy 3 3)
                       )
          }

ply4 :: AI
ply4 = AI { name = "ply4"
          , description = "Minimax with depth 4"
          , strategy =  (ifPieces (==60) 
                         (firstTurn $ minimaxStrategy 4 3)
                         (narrowDoubleCaptures $ minimaxStrategy 4 3)
                        )
          }


dyn1 :: AI
dyn1 = AI { name = "dyn1"
          , description = "Minimax with dynamic depth 1-4"
          , strategy = (ifPieces (==60)
                        (firstTurn greedyStrategy)
                        (ifPieces (>52) 
                         (onlyCaptureStack greedyStrategy)
                         (narrowDoubleCaptures $ 
                          ifPieces (>30)
                          (minimaxStrategy 2 3)
                          (ifPieces (>20)
                           (minimaxStrategy 3 3)
                           (minimaxStrategy 4 5)
                          )
                         )
                        )
                       )
           }


dyn2 :: AI
dyn2 = AI { name = "dyn2"
          , description = "Minimax with dynamic depth 2-6"
          , strategy = (ifPieces (==60)
                        (firstTurn greedyStrategy)
                        (ifPieces (>52) 
                        (onlyCaptureStack $ minimaxStrategy 2 3)
                         (narrowDoubleCaptures $  
                          ifPieces (>30)
                          (minimaxStrategy 3 3)
                          (ifPieces (>20)
                           (minimaxStrategy 4 3)
                           (minimaxStrategy 6 5)
                          )
                         )
                        )
                       )
          }


-- | A greedy strategy: locally maximize the static evaluation function
greedyStrategy :: Strategy
greedyStrategy (GameTree _ branches) rndgen 
    = trace ("Greedy score: " ++ show bestscore) (bestmove, rndgen)
    where 
      choices = [(m, eval $ root t) | (m,t)<-branches]
      (bestmove,bestscore)= maximumBy (\x y -> compare (snd x) (snd y)) choices
      root (GameTree x _) = x




-- | Minimaxing strategy to ply depth `n'
--   With alpha-beta and depth prunning
minimaxStrategy :: Int -> Int -> Strategy
minimaxStrategy n m g rndgen 
    = trace ("Minimax score: " ++ show bestscore) (bestmove, rndgen)
    where g'  = prunebreadth m $  -- ^ cut to breadth $m$
                highfirst $       -- ^ order moves using static evaluation
                mapTree eval $    -- ^ apply evaluation function
                prunedepth n g    -- ^ cut to depth $n$
          (bestmove,bestscore) = minimaxMove_ab (-inf) inf g'


-- | Naive minimax algorithm
--   nodes should contain the static evaluation scores
minimax :: (Num a, Ord a) => GameTree a m -> a 
minimax (GameTree x []) = x
minimax (GameTree _ branches) = - minimum (map (minimax.snd) branches)


-- auxiliary function that returns the best first move
minimaxMove :: (Num a, Ord a) => GameTree a m -> (m,a)
minimaxMove (GameTree _ branches) = (m, -x)
    where (m,x) = minimumBy (\x y ->compare (snd x) (snd y)) [(m,minimax t) | (m,t)<-branches]


-- | Minimax with alpha-beta prunning
minimax_ab :: (Num a, Ord a) => a -> a -> GameTree a m -> a
minimax_ab a b (GameTree x []) = a `max` x `min` b
minimax_ab a b (GameTree _ branches) = cmx a b (map snd branches)
    where cmx a b []  = a
          cmx a b (t:ts) | a'>=b = b
                         | otherwise = cmx a' b ts
                         where a' = - (minimax_ab (-b) (-a) t)


-- | This variant also returns the best move
--   should always be called with a non-empty tree
minimaxMove_ab :: (Num a, Ord a) => a -> a -> GameTree a m -> (m,a)
minimaxMove_ab a b (GameTree x []) = (undefined, a`max`x`min`b)
minimaxMove_ab a b (GameTree _ branches@((m,_):_)) = cmx m a b branches
    where cmx m a b []  = (m,a)
          cmx m a b ((m',t):branches) 
              | a'>=b = (m,b)
              | otherwise = cmx m' a' b branches
              where a' = - (minimax_ab (-b) (-a) t)

      


-- | Static evaluation function
eval :: (Bool,Board) -> Int
eval (True, b) = value b
eval (False,b) = - value (swapBoard b)

value :: Board -> Int
value b@(you,other)
    | pieces==0  || null captures  = -inf
    | pieces'==0 || null captures' = inf
    | otherwise = threats + positional 
    where pieces = length $ nub $ map fst $ Map.elems you
          pieces'= length $ nub $ map fst $ Map.elems you
          captures = nextCaptureMoves b             -- my captures
          captures'= nextCaptureMoves (swapBoard b) -- opponents's captures
          -- the zones of control for each player
          -- the active play has advantage for equal heights
          zoc = zoneOfControl (>=) b
          zoc'= zoneOfControl (>) (swapBoard b)

          -- the three piece types
          ts = [Tzaar, Tzarra, Tott]

          -- immediate threats
          threats = points safe' - points safe 

          -- pieces  safe from immediate threat
          safe = minimum [count t you - min 2 (count t zoc') | t <- ts]
          safe'= minimum [count t other - min 2 (count t zoc) | t <- ts]

          points n | n<=0      = inf`div`2
                   | n==1      = inf`div`4
                   | otherwise = 0

          -- positional score
          -- sum heights multiplied by "relevance" factor 
          -- inside other player's ZoC
          positional = sum [material t zoc * relevance t other | t<-ts] -
                       sum [material t zoc'* relevance t you | t<-ts] 

          -- lower count pieces types are more relevant
          relevance t r = 2^(15-count t r)





-- | count pieces of a particular type
count :: Type -> HalfBoard -> Int
count t r = Map.size $ Map.filter (\(t',_)->t'==t) r

-- | material score by piece type
--   sum height for stacks 
material :: Type -> HalfBoard -> Int
material t r = Map.fold (\(t',h) s->if t==t' then s+h else s) 0 r


-- | The "zone of control" of a player 
-- | the opponent's pieces that can be captured in a turn

zoneOfControl :: (Int->Int->Bool) -> Board  -> HalfBoard
zoneOfControl gt board@(you,other) 
    = Map.filterWithKey forPiece other
    where
      forPiece :: Position -> Piece -> Bool
      forPiece p (_, i) = or $ map (downLine i) $ sixLines p
          where
            downLine, downLine' :: Int -> [Position] -> Bool

            downLine i [] = False
            downLine i (p:ps) 
                = case atPosition board p of
                    Nothing -> downLine i ps
                    Just (True, (_, h)) -> h`gt`i
                    Just (False, (_, j)) -> 
                        or $ map (downLine' (max i j)) $ sixLines p

            downLine' i [] = False
            downLine' i (p:ps) 
                = case atPosition board p of
                    Nothing -> downLine' i ps
                    Just (True, (_, h)) -> h`gt`i
                    Just (False, _) -> False

                                        



-- a higher value than legitimate evaluation scores
inf :: Int
inf = 2^20


-- | narrow the search space: single capture first move
firstTurn :: Strategy -> Strategy
firstTurn s (GameTree node branches) rndgen 
    = s (GameTree node branches') rndgen
    where branches' = [((m,Nothing),g) | ((m,Nothing), g)<-branches]

-- | narrow the search space: consider only capture-stacking turns
onlyCaptureStack ::  Strategy -> Strategy   
onlyCaptureStack s g rndgen = s (narrowTree g) rndgen
    where
      narrowTree :: BoardTree -> BoardTree
      narrowTree (GameTree node@(b, (you,_)) branches)
          | b = GameTree node [ ((m1,Just m2), narrowTree g) 
                                | ((m1,Just m2), g)<-branches,
                                snd m2 `Map.member` you
                              ]
          | otherwise = GameTree node [ (t, narrowTree g) | (t,g)<-branches ]


-- | eliminate double-captures that lead to the same board
narrowDoubleCaptures :: Strategy -> Strategy
narrowDoubleCaptures s g rndgen = s (nubTree g) rndgen
    where
      nubTree :: BoardTree -> BoardTree
      nubTree (GameTree node branches) 
          = GameTree node $ nubBy equiv [(t, nubTree g) | (t,g)<-branches]
          where
            equiv ((m1,Just m2),_) ((m2', Just m1'),_)
                = fst m1/=fst m2 && m1==m1' && m2==m2'
            equiv _ _ = False
                            

-- | use different strategies depedening on the number of pieces left
ifPieces :: (Int -> Bool) -> Strategy -> Strategy -> Strategy
ifPieces cond s1 s2 g@(GameTree (_,(you,other)) branches) rndgen
    | cond n    = s1 g rndgen   -- use the 1st strategy
    | otherwise = s2 g rndgen   -- use the 2nd strategy
    where
      n = Map.size you + Map.size other