-- | Utilities for AI players.
module AI.Utils
  ( winOrPreventLoss
  , mapTree
  , prunedepth
  , prunebreadth
  , prunebreadth_asc
  -- , highfirst
  -- , lowfirst
  , ifPieces
  , ifBranch
  , singleCaptures
  , dontPass
  ) where


import Board
import Data.List (sortBy, minimumBy)
import qualified Data.Map as Map
import System.Random


-- | some auxiliary functions over game trees
-- apply a function to each node
mapTree :: (a->b) -> GameTree a m -> GameTree b m
mapTree f (GameTree x branches) 
    = GameTree (f x) [(m,mapTree f t) | (m,t)<-branches]

-- apply a function to each edge
mapTree' :: (a->b) -> GameTree s a -> GameTree s b
mapTree' f (GameTree x branches) 
    = GameTree x [(f m,mapTree' f t) | (m,t)<-branches]




-- heuristic to order subtrees with highest/lowest values first
highfirst, lowfirst  :: GameTree Int m -> GameTree Int m
highfirst (GameTree x branches) 
    = GameTree x $ sortBy cmp [(m, lowfirst t) | (m,t)<-branches]
    where cmp (_,x) (_,y) = compare (value y) (value x)
          value (GameTree n _)  = n

lowfirst (GameTree x branches) 
    = GameTree x $ sortBy cmp [(m,highfirst t) | (m,t)<-branches]
    where cmp (_,x) (_, y) = compare (value x) (value y)
          value (GameTree n _)  = n



-- prune to a fixed depth
prunedepth :: Int -> GameTree a m -> GameTree a m
prunedepth 0     (GameTree x branches)  = GameTree x []
prunedepth (n+1) (GameTree x branches) 
    = GameTree x [(m,prunedepth n t) | (m,t)<-branches]

-- prune to a fixed breadth
prunebreadth :: Int -> GameTree a m -> GameTree a m
prunebreadth k (GameTree node branches) 
    = GameTree node [(m, prunebreadth k t) | (m,t)<-take k branches]

-- prune to a fixed breadth, ordering nodes by ascending static evalution
prunebreadth_asc :: Ord s => Int -> GameTree s m -> GameTree s m
prunebreadth_asc k (GameTree node branches) 
    = GameTree node branches'
    where 
      branches' = take k $ 
                  sortBy cmp [(m,prunebreadth_asc k t) | (m,t)<-branches]
      cmp (_,x) (_, y) = compare (value x) (value y)
      value (GameTree n _)  = n


  



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

-- | use different strategies dependening on the branching factor
ifBranch :: (Int->Bool) -> Strategy -> Strategy -> Strategy
ifBranch p s1 s2 g@(GameTree (_,(you,other)) branches) rndgen
    | p (length branches)  = s1 g rndgen   -- 1st strategy
    | otherwise            = s2 g rndgen   -- 2nd strategy



-- | Searches BoardTree to a depth of 1 looking for a 
-- | guaranteed win or a preventable loss.
winOrPreventLoss :: Strategy -> Strategy
winOrPreventLoss s (GameTree a branches) = s $ GameTree a branches2
  where
  winning = [ (t, b) | (t, b@(GameTree _ [])) <- branches ]
{-
  losing  = [ t | (t, (GameTree _ branches')) <- branches, 
                           (_, (GameTree _ [])) <- branches' ]
-}
  branches1 = (if not (null winning) 
               then [head winning]
               else if length branches<cutoff 
                    then [ (t,b) | (t,b)<-branches, not_losing b] 
                    else branches)
  branches2 = if null branches1 then [head branches] else branches1
  not_losing (GameTree _ branches) 
      = null [t | (t, GameTree _ []) <- branches]
  cutoff = 100 -- braching upper bound for searching losing moves



-- | narrow the search space: don't consider double-capture or pass moves
singleCaptures ::  Strategy -> Strategy   
singleCaptures s g@(GameTree _ branches) rndgen 
    | null branches' = s g rndgen
    | otherwise      = s g' rndgen
    where
      g'@(GameTree _ branches') = narrow g
      narrow :: BoardTree -> BoardTree
      narrow (GameTree node@(_, (you,_)) branches)
          = GameTree node [ (t, narrow g) 
                            | (t@(_,Just(_,dest)),g)<-branches, 
                            dest`Map.member`you]

-- don't consider pass moves 
dontPass :: Strategy -> Strategy
dontPass s g rndgen = s (narrow g) rndgen
    where
      narrow :: BoardTree -> BoardTree
      narrow (GameTree node branches)
          = GameTree node [ (t, narrow g) | (t@(m1,Just m2),g)<-branches ]