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


-- A greedy strategy
-- chooses the move with highest static evaluation score
greedy :: AI
greedy = AI { name = "greedy"
            , description = "Maximize the static evaluation function"
            , strategy = (ifPieces (==60) 
                          (firstTurn greedyStrategy)
                          (ifPieces (>48)
                           (onlyCaptureStack greedyStrategy)
                           (narrowDoubleCaptures greedyStrategy)
                          )
                         )
            }



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



-- straight minimaxing strategies with increasing depth
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)
                        )
          }

-- dynamic strategies:
-- increase the maximax depth and breadth towards the end game
dyn1 :: AI
dyn1 = AI { name = "dyn1"
          , description = "Minimax with dynamic depth 1-4"
          , strategy = (ifPieces (==60)
                        (firstTurn greedyStrategy)
                        (ifPieces (>48) 
                         (onlyCaptureStack greedyStrategy)
                         (narrowDoubleCaptures $ 
                          ifPieces (>28)
                          (minimaxStrategy 2 3)
                          (ifPieces (>20)
                           (minimaxStrategy 3 4)
                           (minimaxStrategy 4 6)
                          )
                         )
                        )
                       )
           }


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




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


-- Naive minimax algorithm (not used)
-- 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 initial move
minimaxMove_ab :: (Num a, Ord a) => m -> a -> a -> GameTree a m -> (m,a)
minimaxMove_ab m0 a b (GameTree x []) = (m0, a`max`x`min`b)
minimaxMove_ab m0 a b (GameTree _ branches) = cmx m0 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 for a board position
-- boolean indicates if active player is conducting the analysis
eval :: (Bool,Board) -> Int
eval (True, b) = value b
eval (False,b) = - value (swapBoard b)


-- value of a board position for the active player
value :: Board -> Int
value b@(active,other)
    | minimum pieces ==0 || null captures  = -inf
    | minimum pieces'==0 || null captures' =  inf
    | otherwise = material + positional + threats 
    where 
      -- piece counts for each player 
      pieces = counts active
      pieces'= counts other

      captures = nextCaptureMoves b
      captures'= nextCaptureMoves (swapBoard b)    

      -- the zones of control for each player
      -- active player has advantage for equal height
      zoc = zoneOfControl (>=) b
      zoc'= zoneOfControl (>) (swapBoard b)

      -- capture counts by piece type
      nzoc = counts zoc
      nzoc'= counts zoc'

      -- material score
      material = sumHeights active - sumHeights other

      -- positional score
      positional = sumHeights zoc - sumHeights zoc'

      -- scores for immediate threats
      threats = penalty p - penalty q

      p = minimum [x-min 2 y | (x,y)<-zip pieces' nzoc]
      q = minimum [x-min 2 y | (x,y)<-zip pieces nzoc']

      penalty n | n<=2      = inf`div`(2^(1+n))
                | otherwise = 0




-- a higher value than legitimate evaluation score
inf :: Int
inf = 2^10

            
-- count the number of pieces of each type
-- results ordered by piece types 
counts :: HalfBoard -> [Int]
counts b = Map.elems $ Map.fold (\(t,_)-> Map.adjust (+1) t) zeroPieces b

-- finite map assigning 0 to each piece type
-- lifted to top-level to allow sharing across multiple calls
zeroPieces :: Map Type Int
zeroPieces = Map.fromList [(Tzaar,0),(Tzarra,0),(Tott,0)] 

-- sum the heights of pieces (material value of a player)
sumHeights :: HalfBoard -> Int
sumHeights b = sum [h | (_,h)<-Map.elems b]



-- Estimate the "zone of control" of the active player
-- i.e. the opponent's pieces reachable in one or two captures
zoneOfControl ::  (Int->Int->Bool) -> Board -> HalfBoard
zoneOfControl cmp board@(_,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`cmp`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`cmp`i
                    Just (False, _) -> False

                                        




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