module Exercise08 where

import Data.Bits
import Data.List
import System.Random (Random, mkStdGen, randomIO, randoms)

-- 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 :: Size
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 -}

-- helper functions
(.+.) :: Pos -> Pos -> Pos
(.+.) (a, b) (x, y) = (a + x, b + y)

allCells :: Board -> [Pos]
allCells board = [(y, x) | y <- [0 .. height board - 1], x <- [0 .. width board - 1]]

-- Check if the given player can put an orb on the given position
canPlaceOrb :: Player -> Pos -> Board -> Bool
canPlaceOrb player pos board = - player /= signum (getCell pos board)

-- 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 player board = and [canPlaceOrb player (y, x) board | y <- [0 .. height board - 1], x <- [0 .. width board - 1]]

-- the list of neighbors of a cell
neighbors :: Size -> Pos -> [Pos]
neighbors size pos = [pos .+. delta | delta <- [(1, 0), (-1, 0), (0, 1), (0, -1)], isValidPos size (pos .+. delta)]

-- update a single position on the board
updatePos :: (Int -> Int) -> Player -> Pos -> Board -> Board
updatePos f player pos board =
  [ [ if (y, x) == pos
        then f (abs (getCell (y, x) board)) * player
        else getCell (y, x) board
      | x <- [0 .. width board - 1]
    ]
    | y <- [0 .. height board - 1]
  ]

{- x.2 -}

-- place an orb for the given player in the given cell
putOrb :: Player -> Pos -> Board -> Board
putOrb player pos board = infectFilled player (updatePos (+ 1) player pos board)

infectFilled :: Player -> Board -> Board
infectFilled player board
  | hasWon player board = board
  | not (null filled) = infectFilled player (foldr applyOverflow board filled)
  | otherwise = board
  where
    neighborList = neighbors (size board)
    filled = [(y, x) | y <- [0 .. height board - 1], x <- [0 .. width board - 1], abs (getCell (y, x) board) >= length (neighborList (y, x))]
    applyOverflow pos b = foldr (updatePos (+ 1) player) (updatePos (\x -> x - length (neighborList pos)) player pos b) (neighborList pos)

{- x.3 -}

{-WETT-}
-- Heuristic board evaluators
-- this heuristic is taken from https://brilliant.org/discussions/thread/artificial-intelligence-for-chain-reaction/
heuristic :: Player -> Board -> Int
heuristic team board
  | hasWon team board = inf
  | hasWon (- team) board = - inf
  | otherwise = sum (map cellScore (allCells board)) + comboScore
  where
    cellScore :: Pos -> Int
    cellScore pos
      | signum (getCell pos board) == team =
        abs (getCell pos board) +
        let numThreats = length (filter (\n -> isCriticalTeam n board (- team)) (neighbors (size board) pos)) in
        if numThreats > 0
          then numThreats * (maxOrbs pos board - 5)
          else
            case maxOrbs pos board of -- positional score
              2 -> 3  -- corner
              3 -> 2  -- edge
              _ -> 0
              +
              if isCriticalTeam pos board team  -- critical score
                then 2  -- crit
                else 0  -- non-crit
      | otherwise = 0

    comboScore = length (filter isCriticalCombo (allCells board))
    isCriticalCombo p = isCriticalTeam p board team && any (\n -> isCriticalTeam n board team) (neighbors (size board) p)

-- Heuristic helpers
inf :: Int
inf = 100000

isCritical :: Pos -> Board -> Bool
isCritical pos board = abs (getCell pos board) == maxOrbs pos board - 1

isCriticalTeam :: Pos -> Board -> Player -> Bool
isCriticalTeam pos board team = getCell pos board == (maxOrbs pos board - 1) * team

maxOrbs :: Pos -> Board -> Int
maxOrbs pos board = length (neighbors (size board) pos)

-- Your strategy
type Move = (Int, Pos) -- score of the move and the position to place orb

abStrategy :: StatefulStrategyFunc Int
abStrategy turn rng player board = (snd $ alphabeta turn board 2 (- inf) inf player, succ turn)
  where
    alphabeta :: Int -> Board -> Int -> Int -> Int -> Player -> Move
    alphabeta cTurn boardState depth alpha beta team
      | cTurn >= 2 && hasWon player boardState = (inf, (0, 0))
      | cTurn >= 2 && hasWon (- player) boardState = (- inf, (0, 0))
      | depth == 0 = (heuristic player boardState, (0, 0))
      | otherwise =
        if team == player
          then maxVal candidateMoves (- inf, head candidateMoves) alpha beta
          else minVal candidateMoves (inf, head candidateMoves) alpha beta
      where
        candidateMoves = filter (\p -> canPlaceOrb team p boardState) (allCells boardState)

        maxVal [] best a b = best
        maxVal (move : moves) (bestVal, bestMove) a b = if newA >= b then newBest else maxVal moves newBest newA b
          where
            (newVal, _) = alphabeta (succ cTurn) (putOrb team move boardState) (depth - 1) a b (- team)
            newBest = if newVal > bestVal then (newVal, move) else (bestVal, bestMove)
            newA = max a (fst newBest)

        minVal [] worst a b = worst
        minVal (move : moves) (worstVal, worstMove) a b = if newB <= a then newWorst else minVal moves newWorst a newB
          where
            (newVal, _) = alphabeta (succ cTurn) (putOrb team move boardState) (depth - 1) a b (- team)
            newWorst = if newVal < worstVal then (newVal, move) else (worstVal, worstMove)
            newB = min b (fst newWorst)

strategy :: Strategy
strategy rng team board =
  snd $
    foldr
      ( \move (maxScore, maxMove) ->
          let cScore = scoreOfMove move
           in if cScore > maxScore
                then (cScore, move)
                else (maxScore, maxMove)
      )
      (scoreOfMove (head viableFields), head viableFields)
      (tail viableFields)
  where
    scoreOfMove pos = heuristic team (putOrb team pos board)
    viableFields = filter (\p -> canPlaceOrb team p board) (allCells board)

naiveStrategy :: Strategy
naiveStrategy rng team board = foldr (\(y, x) (minY, minX) -> if x + y < minX + minY then (y, x) else (minY, minX)) (head viableFields) viableFields
  where
    viableFields = filter (\p -> canPlaceOrb team p board) (allCells board)

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

-- 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 Int
strategyState = (0, abStrategy)

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