module Exercise08 where

import Data.Bits
import Data.List
import Data.Ord (comparing, compare)
import Control.Monad (ap)
import System.Random (mkStdGen, randoms, randomIO, Random)

-- Player is either 1 or -1
type Player = Int

-- A field is just an Int value where the absolute gives the number of pieces on the field
-- and the sign corresponds to the player
-- e.g. -3 would mean there are three blobs in this field of player -1
type Field = Int

type Row = [Field]
type Column = [Field]
-- boards are rectangles represented as a list of rows
type Board = [Row]

-- A position on the board is represented as (row, column)
-- (0,0) is the top left corner, coordinate values increase towards the bottom right
type Pos = (Int, Int)

-- A size represented as (height,width)
type Size = (Int, Int)

-- A strategy takes the player who's move it is, optionally takes a list of double values
-- to allow for probabilistic strategies, takes the current board and gives back the position
-- of the move the player should do
type Strategy = [Double] -> Player -> Board -> Pos
-- A stateful strategy can additionally pass some object between invocations
type StatefulStrategyFunc a = a -> [Double] -> Player -> Board -> (Pos, a)
-- first value is the state object to pass to the first invocation of each game
type StatefulStrategy a = (a, StatefulStrategyFunc a)

defaultSize :: (Int, Int)
defaultSize = (9,6)

-- Some useful helper functions
row :: Board -> Int -> Row
row = (!!)

column :: Board -> Int -> Column
column = row . transpose

width :: Board -> Int
width (x : _) = length x
width _       = 0

height :: Board -> Int
height = length

size :: Board -> Size
size b = (height b, width b)

getCell :: Pos -> Board -> Field
getCell (y, x) b = b !! y !! x

-- pretty print a single cell
showCell :: Field -> String
showCell c = "- +" !! succ (signum c) : show (abs c)

-- pretty print the given board
showBoard :: Board -> String
showBoard = unlines . map (unwords . map showCell)

-- print a board to the console
printBoard :: Board -> IO ()
printBoard = putStr . showBoard

-- check if a position is one a board of the given size
isValidPos :: Size -> Pos -> Bool
isValidPos (r, c) (y, x) = y >= 0 && y < r && x >= 0 && x < c

{- x.1 -}

-- Check if the given player can put an orb on the given position
canPlaceOrb :: Player -> Pos -> Board -> Bool
canPlaceOrb p c b = (getCell c b * p) >= 0

-- Check if the given player has won the game,
-- you can assume that the opponent has made at least one move before
hasWon :: Player -> Board -> Bool
hasWon p b = all ((>= 0) . (* p)) $ concat b

-- the list of neighbors of a cell
neighbors :: Size -> Pos -> [Pos]
neighbors s (y, x) = filter (inside s) [(y, x + 1), (y, x - 1), (y + 1, x), (y - 1, x)]

inside :: Size -> Pos -> Bool
inside (r, c) (y, x) = 0 <= y && y < r && 0 <= x && x < c

-- update a single position on the board
-- f: function that modifies the number of orbs in the cell
-- p: player to whom the updated cell should belong
updatePos :: (Int -> Int) -> Player -> Pos -> Board -> Board
updatePos f p c b = replace c (f (abs $ getCell c b) `setOwner` p) b

owns :: Player -> Field -> Bool
owns p c = p == signum c

setOwner :: Field -> Player -> Field
setOwner 0 p = 0
setOwner c p
    | p `owns` c = c
    | otherwise  = -c

replace :: Pos -> Int -> Board -> Board
replace (y, x) v b = let (lr, r:rr) = splitAt y b
                         (lc, _:rc) = splitAt x r
                     in lr ++ ((lc ++ (v:rc)):rr)

{- x.2 -}

-- place an orb for the given player in the given cell
putOrb :: Player -> Pos -> Board -> Board
putOrb p c b = putOrb' (size b) p c b 

putOrb' :: Size -> Player -> Pos -> Board -> Board
putOrb' s p c b = let b' = updatePos (+1) p c b in if filled s c b' then overflow s p c b' else b'

filled :: Size -> Pos -> Board -> Bool
filled s c b = criticalSize s c <= abs (getCell c b)

overflow :: Size -> Player -> Pos -> Board -> Board
overflow s p c b = let b' = updatePos (subtract $ criticalSize s c) p c b
                       bb = foldr (putOrb p) b' $ neighbors s c
                   in if hasWon p b then b else bb 

criticalSize :: Size -> Pos -> Int
criticalSize _ (0, 0)      = 2
criticalSize (r, c) (y, x)
    | y == (r-1) && x == 0 = 2
    | (y == 0 || y == (r-1)) && x == (c-1) = 2
criticalSize _ (0, _) = 3
criticalSize _ (_, 0) = 3
criticalSize (r, c) (y, x)
    | y == (r-1) || x == (c-1) = 3
criticalSize _ _ = 4

{- x.3 -}

{-WETT-}
type Score = Int
winScore :: Score
winScore = 10000

-- Score Board Heuristic (from Brilliant)
scoreBoard :: Size -> Player -> Board -> Score
scoreBoard s p b
    | hasWon p b = winScore
    | otherwise  = let os = abs $ sum $ filter ((p ==) . signum) $ concat b
                       cs = sum [ scoreCell s p (y, x) b | y <- [0..fst s - 1], x <- [0..snd s - 1] ]
                       ch = scoreChain s p b
                   in os + cs + ch

ownsCritical :: Size -> Player -> Pos -> Board -> Bool
ownsCritical s p c b = p `owns` getCell c b && isCritical s c b

-- Get the score of the chains of critical cells
scoreChain :: Size -> Player -> Board -> Score
scoreChain s p b = let cs = filter (flip (ownsCritical s p) b) $ cells s
                   in  2 * length (inChain s cs)

-- Get all cells in a chain
inChain :: Size -> [Pos] -> [Pos]
inChain = flip inChain' []

inChain' :: Size -> [Pos] -> [Pos] -> [Pos]
inChain' s xs [] = xs
inChain' s xs (y:ys)
    | any (`elem` (xs ++ ys)) (neighbors s y) = inChain' s (y:xs) ys
    | otherwise = inChain' s xs ys

-- Generate all possible cells
cells :: Size -> [Pos]
cells s = [(y, x) | y <- [0..fst s - 1], x <- [0..snd s - 1]]

isCritical :: Size -> Pos -> Board -> Bool
isCritical s c b = criticalSize s c - 1 == abs (getCell c b)

-- Score a single cell
scoreCell :: Size -> Player -> Pos -> Board -> Score
scoreCell s p c b
    | p `owns` getCell c b = let ec = filter (flip (ownsCritical s p) b) $ neighbors s c
                                 es = if null ec then scorePos s c + scoreCritical s c b
                                      else length ec * (criticalSize s c - 5)
                             in es -- * abs (getCell c b)
    | otherwise = 0

scorePos :: Size -> Pos -> Score
scorePos s c
    | criticalSize s c < 4 = 5 - criticalSize s c
    | otherwise = 0

scoreCritical :: Size -> Pos -> Board -> Score
scoreCritical s c b
    | isCritical s c b = 2
    | otherwise = 0

type Depth = Int

-- Generate all possible moves
moves :: Size -> [Pos] -> Player -> Board -> [Pos]
moves s cs p b = filter (not . owns (-p) . flip getCell b) cs

type ScoredMove = (Move, Score, Board)
type Move = Pos

scoredMoves :: (Move -> ScoredMove) -> [Move] -> [ScoredMove]
scoredMoves = map 

topN :: Int -> [ScoredMove] -> [ScoredMove]
topN 0 ms = sortOn getScore ms
topN n ms = take n (sortOn getScore ms)

scoreMove :: Size -> Player -> Pos -> Board -> Score
scoreMove s p c b = scoreBoard s p $ putOrb' s p c b

scoreMoveBoard :: Size -> Player -> Pos -> Board -> ScoredMove
scoreMoveBoard s p c b = let b' = putOrb' s p c b
                         in (c, scoreBoard s p b', b')

-- Get the Best Move
bestMove :: Size -> [Pos] -> Player -> Board -> Pos
bestMove s cs p b = maximumBy (comparing (flip (scoreMove s p) b)) $ moves s cs p b

bestScore :: Size -> [Pos] -> Player -> Board -> Score
bestScore s cs p b = maximum $ map (flip (scoreMove s p) b) $ moves s cs p b

-- Get the best Move Score Pair
bestMoveScore :: Size -> [Pos] -> Player -> Board -> (Move, Score)
bestMoveScore s cs p b = maximumBy (comparing snd) $ map (ap (,) $ flip (scoreMove s p) b) $ moves s cs p b

getScore :: ScoredMove -> Score
getScore (_, s, _) = s

data TreeData = T Board Size [Pos]
generateTreeData :: Board -> TreeData
generateTreeData b = let s = size b in T b s $ cells s

getBoard :: TreeData -> Board
getBoard (T b _ _) = b

getSize :: TreeData -> Size
getSize (T _ s _) = s

getCells :: TreeData -> [Pos]
getCells (T _ _ cs) = cs

setBoard :: TreeData -> Board -> TreeData
setBoard (T _ s cs) b = T b s cs

doMoves :: TreeData -> Player -> [(Move, Board)]
doMoves (T b s cs) p = map (\m -> (m, putOrb' s p m b)) (moves s cs p b)

-- Alpha Beta Prunning
type Alpha = Score
type Beta  = Score
abPrunning :: Depth -> TreeData -> Alpha -> Beta -> Player -> (Move, Score)
-- abPrunning 0 (T b s _) _ _ p = ((0,0), (-p) * scoreBoard s (-p) b) -- + p * scoreBoard s p b
abPrunning 1 d _ _ p   = abPrunningTerminal d p
abPrunning n d a' b' p = let ms = doMoves d p
                             sc = abPrunning' (n - 1) d a' b' p ms ((0, 0), (-p) * winScore)
                         in sc

abPrunningTerminal :: TreeData -> Player -> (Move, Score)
abPrunningTerminal (T b s cs) p = let (m, sc) = bestMoveScore s cs p b
                                  in (m, p * sc - (-p) * scoreBoard s (-p) (putOrb' s p m b))

abPrunning' :: Depth -> TreeData -> Alpha -> Beta -> Player -> [(Move, Board)] -> (Move, Score) -> (Move, Score)
abPrunning' _ _ a' b' _ _ mv | b' <= a'            = mv
abPrunning' _ _ _ _ _ [] mv                        = mv
abPrunning' _ _ _ _ p ((m, b):_) _ | hasWon p b    = (m, p * winScore)
abPrunning' n d a' b' (-1) ((m, b):ms) mv@(_, sc') = let (_, sc) = abPrunning n (d `setBoard` b) a' b' 1
                                                         new_b   = min b' sc
                                                         new_mv  = if sc' < sc then mv else (m, sc)
                                                     in abPrunning' n d a' new_b (-1) ms new_mv
abPrunning' n d a' b' 1 ((m, b):ms) mv@(_, sc')    = let (_, sc) = abPrunning n (d `setBoard` b) a' b' (-1)
                                                         new_a   = max a' sc
                                                         new_mv  = if sc' > sc then mv else (m, sc)
                                                     in abPrunning' n d new_a b' 1 ms new_mv


abPrunningStrat :: StatefulStrategyFunc (Maybe (Int, TreeData))
abPrunningStrat Nothing x p b = abPrunningStrat (Just (0, generateTreeData b)) x p b
abPrunningStrat dt@(Just (n, t)) _ p b
    | n < 2     = (bestMove (getSize t) (getCells t) p b, Just (n + 1, t))
    | otherwise = let d = t `setBoard` b in (fst $ abPrunning 3 d (-winScore) winScore p, dt)

-- Min max
minMax :: Depth -> Maybe BoardData -> Player -> Board -> (ScoredMove, Maybe BoardData)
minMax d Nothing p b           = let s = size b in minMax d (Just $! B s $! cells s) p b
minMax 1 d@(Just (B s cs)) p b = let (m, sc) = bestMoveScore s cs p b in ((m, sc, []), d)
minMax n d@(Just (B s cs)) p b = let (m:ms) = topN 0 $! map (flip (scoreMoveBoard s p) b) (moves s cs p b)
                                     bs = maximumBy (comparing getScore) $! map rec (m:ms)
                                 in if getScore m == winScore then (m, d) else (bs, d)
    where
        rec :: ScoredMove -> ScoredMove
        rec (m, sc, b') = let ((_, minMaxScore, _), _) = minMax (n-1) d (-p) b'
                          in (m, sc - minMaxScore, b')


bestMovePredictive :: Size -> [Pos] -> Player -> Board -> Pos
bestMovePredictive s cs p b = maximumBy (comparing scoreFunc) $! moves s cs p b
    where scoreFunc m = let (_, sc, b') = scoreMoveBoard s p m b
                        in if sc == winScore then sc else sc - (bestScore s cs (-p) $! b')

 -- Your strategy
strategy :: Strategy
strategy _ p b = let s = size b in bestMovePredictive s (cells s) p b

-- adds state to a strategy that doesn't use it
wrapStrategy :: Strategy -> StatefulStrategy Int
wrapStrategy strat = (0, \s r p b -> (strat r p b, succ s))

data BoardData = B Size [Pos]

strategy' :: StatefulStrategyFunc (Maybe BoardData)
strategy' d r p b = (strategy r p b, d)
-- strategy' d _ p b = let ((m, _, _), d') = minMax 3 d p b in (m, d')
--strategy' Nothing _ p b  = let s  = size b 
--                               cs = cells s
--                           in (bestMovePredictive s cs p b, Just (B s cs))
-- strategy' (Just (B s cs)) _ p b = (bestMovePredictive s cs p b, Just (B s cs))

-- the actual strategy submissions
-- if you want to use state modify this instead of strategy
-- additionally you may change the Int in this type declaration to any type that is usefully for your strategy
-- strategyState :: StatefulStrategy (Maybe BoardData)
strategyState = (Nothing, strategy') 
-- strategyState = (Nothing, abPrunningStrat) 
{-TTEW-}

-- Simulate a game between two strategies on a board of the given size and
-- returns the state of the board before each move together with the player that won the game
play :: [Int] -> Size -> StatefulStrategy a -> StatefulStrategy b -> [(Board, Pos)]
play rss (r, c) (isa, sa) (isb, sb) = go rss isa sa isb sb 1 0 (replicate r (replicate c 0))
  where
    -- type signature is necessary, inferred type is wrong!
    go :: [Int] -> a -> StatefulStrategyFunc a -> b -> StatefulStrategyFunc b -> Player -> Int -> Board -> [(Board, Pos)]
    go (rs:rss) stc sc stn sn p n b
      | won       = []
      | valid     = (b, m) : go rss stn sn st' sc (-p) (succ n) (putOrb p m b)
      | otherwise = []
      where
        won      = n > 1 && hasWon (-p) b
        (m, st') = sc stc (mkRandoms rs) p b
        valid    = isValidPos (size b) m && canPlaceOrb p m b

-- Play a game and print it to the console
playAndPrint :: Size -> StatefulStrategy a -> StatefulStrategy b -> IO ()
playAndPrint size sa sb = do
  seed <- randomIO
  -- let seed = 42
  let moves = play (mkRandoms seed) size sa sb
  putStr $
    unlines (zipWith showState moves $ cycle ['+', '-']) ++ "\n" ++
    (case length moves `mod` 2 of { 1 -> "Winner: +"; 0 -> "Winner: -" }) ++ "\n" ++
    "View at https://vmnipkow16.in.tum.de/christmas2020/embed.html#i" ++ base64 (1 : t size ++ concatMap (t . snd) moves) ++ "\n"
  where
    showState (b, pos) p = showBoard b ++ p : " places at " ++ show pos ++ "\n"
    t (a, b) = [a, b]

mkRandoms :: Random a => Int -> [a]
mkRandoms = randoms . mkStdGen

base64 :: [Int] -> String
base64 xs = case xs of
  []            -> ""
  [a]           -> f1 a : f2 a 0 : "=="
  [a, b]        -> f1 a : f2 a b : f3 b 0 : "="
  a : b : c : d -> f1 a : f2 a b : f3 b c : f4 c : base64 d
  where
    alphabet = (!!) "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    f1 a     = alphabet $ shiftR a 2
    f2 a b   = alphabet $ shiftL (a .&. 3 ) 4 .|. shiftR b 4
    f3 b c   = alphabet $ shiftL (b .&. 15) 2 .|. shiftR c 6
    f4 c     = alphabet $ c .&. 63