module Exercise08 where

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

-- 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 (y, x) b = p * (b !! y !! x) >= 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 = if p > 0 then all (>= 0) $ concat b else all (<= 0) $ concat b

-- the list of neighbors of a cell
neighbors :: Size -> Pos -> [Pos]
neighbors (r, c) (y, x) = filter (\(y', x') -> y' >= 0 && y' < r && x' >= 0 && x' < c) [(y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)]

-- 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 (y, x) b = let col = b !! y in take y b ++ [take x col ++ [p * f (abs (col !! x))] ++ drop (x + 1) col] ++ drop (y + 1) b

{- x.2 -}

-- place an orb for the given player in the given cell
putOrb :: Player -> Pos -> Board -> Board
putOrb p pos b = decider p ([pos], updatePos (+1) p pos b)

-- stops when the player has won or when the board is stable again
decider :: Player -> ([Pos], Board) -> Board
decider p (pos, b) = if hasWon p b || null overflown then b else decider p updated
  where overflown = getOverflown p pos b
        updated = dealWithOF p overflown ([], b)

-- in: list of potentially overflown positions
-- out: list of actually overflown positions
getOverflown :: Player -> [Pos] -> Board -> [Pos]
getOverflown p pos b = [(y, x) | (y, x) <- nub pos, abs (b !! y !! x) >= 4 - (if y == 0 || y == r - 1 then 1 else 0) - (if x == 0 || x == c - 1 then 1 else 0)]
  where (r, c) = size b

-- in: list of actually overflown positions and current board
-- out: list of new potentially overflown positions and updated board
dealWithOF :: Player -> [Pos] -> ([Pos], Board) -> ([Pos], Board)
dealWithOF _ [] res = res
dealWithOF pl (p:ps) (pos, b) = dealWithOF pl ps (pos ++ ns, updatedBoard)
  where (updatedBoard, ns) = flowOver pl p b

flowOver :: Player -> Pos -> Board -> (Board, [Pos])
flowOver p (y, x) b = (before ++ middle : after, neighbours)
  where (start, rest) = splitAt y b
        end = tail rest 

        before = if null start then [] else init start ++ [start' ++ (p * (abs (head rest') + 1)) : end']
          where (start', rest') = splitAt x $ last start
                end' = tail rest'

        (middle, neighbours) = (firstField start' ++ secondField (head rest') : thirdField end', neighbours)
          where (start', rest') = splitAt x $ head rest
                end' = tail rest'
                neighbours = [(y - 1, x) | not (null start)] ++ [(y + 1, x) | not (null end)] ++
                             [(y, x - 1) | not (null start')] ++ [(y, x + 1) | not (null end')]

                firstField [] = []
                firstField fields = init fields ++ [p * (abs (last fields) + 1)]

                secondField field = p * (abs field - length neighbours)

                thirdField [] = []
                thirdField fields = (p * (abs (head fields) + 1)) : tail fields

        after = if null end then [] else (start' ++ (p * (abs (head rest') + 1)) : end') : tail end
          where (start', rest') = splitAt x $ head end
                end' = tail rest'

{- x.3 -}

{-WETT-}

strategyState :: StatefulStrategy Options
strategyState = (Uncalculated, statefulStrategyFunc) 

-- returns a "random" Element of a list
getRandElem :: Double -> [a] -> a
getRandElem d as = as !! floor (d / 100 * 100 * fromIntegral (length as))

-- ONLY use when the parameter is definitely not null!
unmaybe :: Maybe a -> a
unmaybe (Just a) = a

-- the tree of all the informations
type State = (Board, LastMoveBy, Options)
type LastMoveBy = Player
data Options = Finnished | Calculated [(State, Pos)] | Uncalculated  -- Calculated and Uncalculated both imply unfinnished
  deriving (Show, Eq)
type Score = Int

statefulStrategyFunc :: StatefulStrategyFunc Options
statefulStrategyFunc options ds p b = 
  getRandElem (head ds) $ 
  getBestChoicesWithOptions $ 
  expandLayersRoot 2200 $
  getCurrentState options $ map (map (*p)) b  -- normalises -> we're always Player +1

-- updates the root to what option our opponent chose
-- if there's no root with the current board a new root is initialized
getCurrentState :: Options -> Board -> State
getCurrentState (Calculated c) b = if null found then (b, -1, Uncalculated) else unmaybe found
  where found = find (\(b', _, _) -> b' == b) $ map fst c
getCurrentState _ b = (b, -1, Uncalculated)

getBestChoicesWithOptions :: State -> [(Pos, Options)]
getBestChoicesWithOptions (_, _, Calculated c) = [(pos, o) | (pos, (s, o)) <- choices, s == bestScore]
  where choices = map (\((b, p, o), pos) -> (pos, (score 0 (b, p, o), o))) c 
        bestScore = bestScoreTupel $ map (fst . snd) choices
getBestChoicesWithOptions _ = []

-- returns: (p's score, -p's score)
score :: Int -> State -> (Score, Score)
score n (_, _, Finnished) = (10000 - n, n - 10000)  -- the further down the Finnished state is, the less likely it is we'll actually get there
score n (b, _, Calculated c) = swap $ bestScoreTupel $ map (score (n + 1) . fst) c  -- swap the best scoreTupel for -p to get the worst one for p
score n (b, p, Uncalculated) = (calcScore normalised, calcScore $ map (map negate) normalised)
  where normalised = map (map (*p)) b

-- first priority: the best option for the choosing player
-- second priority: the worst option for the other player
bestScoreTupel :: [(Score, Score)] -> (Score, Score)
bestScoreTupel = minimumBy sortingScoreTupels
  where sortingScoreTupels (a1, b1) (a2, b2)
          | a1 > a2 = LT 
          | a1 < a2 = GT
          | b1 > b2 = GT
          | b1 < b2 = LT 
          | otherwise = EQ

-- criteria inspired by: https://brilliant.org/wiki/chain-reaction-game/
-- from player +1's POV
-- assumes no one has won yet (it won't crash but might not quite work as intended if someone did already win)
calcScore :: Board -> Int
calcScore b = sum fieldScores
  where (r, c) = size b 
        fs = concat b

        -- mappes each field to its mass and the positions of its neighbours
        mappedBoard = [let rw = b !! y in [(rw !! x, neighbors (r, c) (y, x)) | x <- [0..c-1]] | y <- [0..r-1]]
        
        -- creates a list of every captured field: its mass and every neighbour's mass and the positions of its neighbours too
        capturedFields = [(mass, [mappedBoard !! y !! x | (y, x) <- positions]) | (mass, positions) <- concat mappedBoard, mass > 0]

        -- mapps the list to a list of the three necessary variables: the field's current mass, its critical mass, 
        -- and the amount of neighbours that are critical and the enemy's
        necessaryVariables = [(mass, length neighbours, length $ filter (\ (neighbourMass, neighbourNeighbours) -> 1 - neighbourMass == length neighbourNeighbours) neighbours) | (mass, neighbours) <- capturedFields]

        -- calculates each fields score:
        --  amound of orbs
        --  for every captured field: - amount of critical enemy neighbours * (5 - critical mass) 
        --    if no critical enemy neighbours and corner field then: + 3
        --    if no critical enemy neighbours and edge field then: + 2
        --    if no critical enemy neighbours and critical then: + 2
        fieldScores = 
          map (\ (mass, critMass, critEnemyNeigh) -> 
                  mass - critEnemyNeigh * (5 - critMass) +
                    if critEnemyNeigh == 0 then 
                      (if critMass < 4 then 5 - critMass else 0) + 
                      (if critMass == mass + 1 then 2 else 0) 
                    else 0
              ) 
          necessaryVariables

-- decides at the root
expandLayersRoot :: Int -> State -> State
expandLayersRoot 0 state = state
expandLayersRoot _ (b, p, Finnished) = (b, p, Finnished)
expandLayersRoot n (b, p, Calculated c) = (b, p, expandedLayersNode n c)
expandLayersRoot n (b, p, Uncalculated) = expandedLayersLeaf n b (negate p)

-- distributes at the nodes
expandedLayersNode :: Int -> [(State, Pos)] -> Options
expandedLayersNode n c = Calculated $ zip (map (expandLayersRoot (n `div` length c)) states) positions
  where (states, positions) = unzip c

-- expands at the leaves
expandedLayersLeaf :: Int -> Board -> Player -> State
expandedLayersLeaf n b p
  | n' < 1 = (b, p, Calculated cleanedOptions)  -- this was the last extension
  | otherwise = (b, p, Calculated $ zip (map (expandLayersRoot n') states) positions)
  where unwrappedOptions = calculateOptions b p
        n' = n `div` length unwrappedOptions - 1 -- because calculating takes time too
        -- if it's our turn and there is a Finnished option then only that one is needed
        firstFinnished = find (\((_, _, o), _) -> o == Finnished) unwrappedOptions
        cleanedOptions = if p == -1 || null firstFinnished then unwrappedOptions else [unmaybe firstFinnished] 
        (states, positions) = unzip cleanedOptions

-- assumes no one has won yet
calculateOptions :: Board -> Player -> [(State, Pos)]
calculateOptions b p = concat [optionsRow y | y <- [0..r-1]]
  where (r, c) = size b 
        optionsRow y = let row = b !! y in [(optionsField y x, (y, x)) | x <- [0..c-1], p * (row !! x) >= 0]
        optionsField y x = (b', p, options)
          where b' = putOrb p (y, x) b
                options = if (p == 1 && any (< 0) (concat b')) || (p == -1 && any (> 0) (concat b')) then Uncalculated else Finnished 

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