-- | Utilities for computing AI game trees
module AI.Tree
  ( BoardTree
  , GameTree(..)
  , Strategy
  , AI (..)
  , Eval
  , boardTree
    --  , startBoardTree
  , mapTree
  , mapTree'
  , isEmptyTree
  , infinity
  -- , winOrPreventLoss
  , pruneDepth, pruneBreadth
  , highFirst, lowFirst
  , withBoard
  , singleMoves
    --  , dontPass, singleCaptures  
  ) where

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

-- | A game tree with nodes s and moves m
data GameTree s m = GameTree s [(m, GameTree s m)] 
                    deriving Show

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

-- | test for empty branches
isEmptyTree :: GameTree a m -> Bool
isEmptyTree (GameTree _ []) = True
isEmptyTree _               = False


-- | A game tree of boards labeled by moves
type BoardTree = GameTree Board Move

-- | An AI strategy calculates the next turn from a board tree.
--  result: evaluation score, next move, next random generator
type Strategy = BoardTree -> StdGen -> (Int, Move, StdGen)

-- | An AI player.
data AI = AI
  { name        :: String   -- ^ Name of AI.
  , description :: String   -- ^ Brief description of AI.
  , strategy    :: Strategy -- ^ The strategy.
  }

-- | type of static evaluation functions
type Eval = Board -> Int  

-- | maximum absolute value of static evaluation 
infinity :: Int
infinity = 2^20


-- | Create a board tree from a board position
boardTree :: Board -> BoardTree
boardTree b = GameTree b [(m, boardTree (applyMove m b)) | m<-nextMoves b]


-- | order subtrees with ascending or descending order (not used)
highFirst, lowFirst  :: GameTree Int m -> GameTree Int m
highFirst (GameTree x branches) 
    = GameTree x [(m,highFirst t) | (m,t)<-sortBy cmp branches] 
    where cmp (_,x) (_,y) = compare (value y) (value x)
          value (GameTree n _)  = n

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



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


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


  
-- | conditional strategy (depending on the board state)
withBoard :: (Board -> Strategy) -> Strategy
withBoard f t@(GameTree b _) g = f b t g

-- number of stacks of both players
-- withStacks :: (Int -> Strategy) -> Strategy
-- withStacks f = withBoard $ \b -> f (IntMap.size (active b) + IntMap.size (inactive b))

-- | avoid search when only move is available
singleMoves :: Strategy -> Strategy
singleMoves s (GameTree b [(m,_)]) rnd = (0, m, rnd)
singleMoves s bt rnd = s bt rnd


{-
-- | Searches BoardTree to a depth of 1 looking for a 
-- | guaranteed win or a preventable loss.
winOrPreventLoss :: Strategy -> Strategy
winOrPreventLoss s (GameTree node branches) = s $ GameTree node branches2
  where
    -- ensure a win in 1 or 2 captures
    winning = [ (m, t) | (m,t@(GameTree b _))<-branches, endGame b]
    -- prevent a loss
    prevent_loss = [(m1,t1) | (m1,t1@(GameTree _ branches'))<-branches,
                    (m2,t2)<- branches', not_losing t2]
                              
    branches1 = (if not (null winning) 
                 then [head winning]
                 else prevent_loss
                )
    branches2 = if null branches1 then [head branches] else branches1
    not_losing (GameTree _ branches) 
        = null [m | (m, GameTree b _) <- branches, endGame b]
    cutoff = 1000 -- 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 board branches)
          = GameTree board [ (m, narrow g) | (m,g)<-branches,
                                 move board==1 || isStacking m ]
      isStacking (Stack _ _) = True
      isStacking _           = False

-- don't consider passing moves 
dontPass :: Strategy -> Strategy
dontPass s g rndgen = s (narrow g) rndgen
    where
      narrow :: BoardTree -> BoardTree
      narrow (GameTree node branches) 
        | null branches' = GameTree node branches
        | otherwise =  GameTree node branches'
        where branches' = [(m, narrow g) | (m,g)<-branches, m/=Pass]
-}