module Exercise_10 where
import Data.List hiding (find, insert)
import qualified Data.List as L
import Data.Maybe
import qualified Test.QuickCheck as QC
import qualified Control.Exception as E
import System.Environment

{-H10.1-}
data Player = V | H -- vertical or horizontal player
  deriving (Eq,Show)
data Field = P Player | E -- a field is occupied by a player or empty
  deriving (Eq,Show)
type Row = [Field]
type Column = [Field]
type Board = [Row] -- we assume that boards are squares and encode a board row by row
data Game = Game Board Player -- stores the currenty board and player
  deriving (Eq,Show)

-- get a given row of a board
row :: Board -> Int -> Row
row = (!!)

-- get a given column of a board
column :: Board -> Int -> Column
column = row . transpose

-- width of a board
width :: Board -> Int
width [] = 0
width (x:xs) = length x

-- height of a board
height :: Board -> Int
height = length

{-H10.1.1-}
prettyShowBoard :: Board -> String
prettyShowBoard = intercalate "\n" . map (concat . map showField)
  where
    showField E = "+"
    showField (P p) = show p

prettyShowGame :: Game -> String
prettyShowGame g@(Game b p) = "It's " ++ show p ++ "'s turn. The player has a " ++ show (lead g) ++ " point advantage.\n" ++ prettyShowBoard b  ++ "\n" ++ show b

{-H10.1.2-}
-- position on a board (row, column)
-- (0,0) corresponds to the top left corner
type Pos = (Int, Int)

isValidMove :: Game -> Pos -> Bool
isValidMove (Game b H) (r,c) = take 2 (drop c (row b r)) == [E, E]
isValidMove (Game b V) (r,c) = let b' = drop r b in
    r <= 10 && take 1 (drop c (head b')) == [E] && take 1 (drop c (head (tail b'))) == [E]


{-H10.1.3-}
--Optimize later
canMove :: Game -> Bool
canMove =  (>0) . validMoves

{-H10.1.4-}
updateBoard :: Board -> Pos -> Field -> Board
updateBoard b p f' = [ [ if p == (ri,ci) then f' else f | (f,ci) <- zip r' [0..11]] | (r',ri) <- zip b [0..11]]

{-H10.1.5-}
playMove :: Game -> Pos -> Game
playMove (Game brd H) (r,c) =  Game (b ++ row' : b') V
  where
    (b,row:b') = splitAt r brd
    row' = let (a,b) = splitAt c row in a ++ [P H, P H] ++ drop 2 b
playMove (Game brd V) (r,c) =  Game (b ++ r1' : r2' : b') H
  where
    (b,r1:r2:b') = splitAt r brd
    replace row = let (a,b) = splitAt c row in a ++ P V : tail b
    r1' = replace r1
    r2' = replace r2


{-H10.1.6-}
-- the first parameter of a strategy is an infinite list of
-- random values between (0,1) (in case you wanna go wild with
-- probabilistic methods)
type Strategy = [Double] -> Game -> Pos

data Res = Loss | Draw Int | Win deriving (Eq , Show)

instance Ord Res where
  (<=) Loss _ = True
  (<=) (Draw a) (Draw b) = a <= b
  (<=) _ Win = True
  (<=) _ _ = False
  compare Loss Loss = EQ
  compare Loss _ = LT
  compare (Draw a) Loss = GT
  compare (Draw a) (Draw b) = compare a b
  compare (Draw a) Win = LT
  compare Win Win = EQ
  compare Win _ = GT

{-WETT-}

minDistance :: Game -> Pos -> Int
minDistance (Game b p) (x,y) = case [ field_val x' y' | (r,x') <- zip b [0..11], (f,y') <- zip r [0..11], isOp f] of
    [] -> 0
    xs -> minimum xs
    where
      (x2,y2) = case p of
        V -> (succ x, y)
        H -> (x, succ y)
      field_val x' y' = squareRoot (min ((x'-x)^2 + (y'-y)^2) ((x'-x2)^2 + (y'-y2)^2))
      isOp = (==) (P (otherPlayer p))
      (^!) :: Num a => a -> Int -> a
      (^!) x n = x^n
      squareRoot :: Int -> Int
      squareRoot 0 = 0
      squareRoot 1 = 1
      squareRoot n =
        let twopows = iterate (^!2) 2
            (lowerRoot, lowerN) =
                last $ takeWhile ((n>=) . snd) $ zip (1:twopows) twopows
            newtonStep x = div (x + div n x) 2
            iters = iterate newtonStep (squareRoot (div n lowerN) * lowerRoot)
            isRoot r  =  r^!2 <= n && n < (r+1)^!2
        in  head $ dropWhile (not . isRoot) iters

-- Returns the best available moves
maxPrune :: Game -> ([Pos], Int)
maxPrune g = case [ ((r,c),lead (playMove g (r,c))) |  r <- [0..11], c <- [0..11], isValidMove g (r,c) ] of
  [] -> ([],-1)
  moves -> (map fst moves', best)
    where
      best = minimum $ map snd moves
      moves' = filter (\(_,rem) -> rem == best) moves


lead :: Game -> Int
lead g@(Game b p) = validMoves g - validMoves (Game b (otherPlayer p))

{-
Computes the score for a game state. It does not return the actual count of
  remaining moves, as it counts safe moves twice.
-}
validMoves :: Game -> Int
validMoves g@(Game b V) = validMoves (Game (transpose b) H)
validMoves g@(Game b H) = foldl' (\acc (p,c,n) -> acc + eval (zip3 p c n)) 0 (zip3 (replicate 12 (P V) : take 11 b) b (drop 1 b ++ [replicate 12 (P V)]))
  where
    value E E = 2
    value E (P _) = 1
    value (P _) E = 1
    value (P _) (P _) = 0
    safe_factor = 2
    valueAdj ((es,ss),acc) = let (safes,safes_rem) = quotRem ss 2 in acc + safe_factor * safes + (es+safes_rem) `div` 2
    collapse_safes :: ((Int,Int),Int) -> ((Int,Int),Int)
    collapse_safes ((es,ss),acc) = ((carry,0),acc')
      where
        (safes,safes_rem) = quotRem ss 2
        regs = (es+safes_rem) `div` 2
        safes_for_regs = es `mod` 2 == 1
        carry = if safes_for_regs then 1 else ((safes_rem + 1) `mod` 2)
        acc' = acc + safe_factor * safes + regs + if not safes_for_regs then (safes_rem + 1) `div` 2 else 0
    eval :: [(Field, Field, Field)] -> Int
    eval = valueAdj . foldl' (\s@((es,ss),acc) f@(p,c,n) -> case (c,es,value p n,ss) of
      (E,_,0,_) -> ((es,succ ss), acc) --Safe field
      (E,_,v,0) -> ((succ es,0), acc)  --Regular field
      (E,_,v,_) -> collapse_safes s --Contiguous safe over -> carry over into acc
      -- (_,0,_,_) -> s                             --wait for E
      _ -> ((0,0), valueAdj s) --fields over -> reset
      ) ((0,0),0)

choose :: [Double] -> Int -> [a] -> ([a],[Double])
choose ds n xs = case compare n len_xs of
  LT -> let (is,ds') = indices' n len_xs [] ds in (map (xs!!) is,ds')
  EQ -> (xs,ds)
  GT -> error "choose too much"
  where
    len_xs = length xs
    indices' :: Int -> Int -> [Int] -> [Double] -> ([Int],[Double])
    indices' 0 _ is ds = (is,ds)
    indices' n o is (d:ds) = if find i_candidate is
      then indices' n o is ds
      else indices' (pred n) o (insert i_candidate [] is) ds
      where
        i_candidate =  (round . (*(fromIntegral $ pred o))) d
        find _ [] = False
        find i (x:xs) = case compare i x of
          LT -> False
          EQ -> True
          GT -> find i xs
        insert i as [] = reverse (i : as)
        insert i as r@(b:bs) = case compare i b of
          LT -> reverse as ++ (i:r)
          EQ -> reverse as ++ r
          GT -> insert i (b:as) bs

--Prefer moves that are closer to the enemy
mkcds :: Int -> [Double] -> Game -> ([Pos], [Double])
mkcds trim ds g = make (s_cds) index_count
  where
    cds = (fst $ maxPrune g)
    len_cdss = length cds
    index_count = max 0 (min trim len_cdss)
    s_cds = groupBy (\(_,a) (_,b) -> a == b) $ sortOn snd $ map (\p -> (p,minDistance g p)) cds
    make [] _ = ([],ds)
    make (cds:cdss) 0 = ([],ds)
    make (cds:cdss) i = let l = length cds in if i > l
      then let (ps, ds') = make cdss (i-l) in ((map fst cds)++ps,ds')
      else choose ds i (map fst cds)


christmasAI :: Strategy
christmasAI = treeStrategy 4 4

data GameTree = Node Res [(Pos, GameTree)] deriving Eq

treeStrategy :: Int -> Int -> Strategy
treeStrategy d w ds g = fst $ fromJust $ L.find (\(_,(Node t_h' _)) -> t_h == t_h') moves
  where
    ds' = drop 1 ds
    close_optimals = map fst $ head $ groupBy (\(_,a) (_,b) -> a == b) $ sortOn snd $ map (\p -> (p,minDistance g p)) optimals
    optimals = map fst $ filter (\(p,(Node t_h' _)) -> t_h == t_h') moves
    Node _ moves = christmasTree ds' d w g
    mx@(p,t@(Node t_h _)) = maximumBy (\(_,(Node a _)) (_,(Node b _)) -> compare a b) moves

-- receives a game and plays a move for the next player
christmasTree :: [Double] -> Int -> Int -> Game -> GameTree
christmasTree ds depth width = play depth ds
  where
    me n = (depth - n) `mod` 2 == 0
    play :: Int -> [Double] -> Game -> GameTree
    play 0 _ g = Node (Draw ((*(if me 0 then 1 else -1))(lead g))) []
    play n ds g = case mkcds width ds g of
      ([],_) -> Node (if me n then Win else Loss) []
      (cs,ds') -> let ts = map (play (pred n) ds' . playMove g) cs
        in Node (reduce ts)  (zip cs ts)
          where
            reduce :: [GameTree] -> Res
            reduce = foldl' (\acc (Node r _) -> red acc r) (if me n then Loss else Win)
              where
                red = if me n then max else min

{-TTEW-}

{-H10.1.7-}
otherPlayer :: Player -> Player
otherPlayer V = H
otherPlayer H = V
emptyBoard :: Int -> Board
emptyBoard dim = replicate dim (replicate dim E)
play :: [[Double]] -> Int -> Strategy -> Strategy -> ([Board],Player)
play rss dim sv sh = go rss (Game (emptyBoard dim) V) []
  where
    chooseStrat H = sh
    chooseStrat V = sv
    go (rs:rss) g@(Game _ p) bs
      | not (canMove g) = gameOver
      | otherwise = let sp = chooseStrat p rs g in
        if isValidMove g sp then
          let g'@(Game nb _) = playMove g sp in
          go rss g' (nb:bs)
        else gameOver
      where gameOver = (reverse bs, otherPlayer p)

-- generates infinite list of values between (0,1)
genRandomZeroOne :: QC.Gen [Double]
genRandomZeroOne = mapM (const $ QC.choose (0::Double,1)) [1..]

-- plays a game and prints it to the console
playAndPrint :: Int -> Strategy -> Strategy -> IO ()
playAndPrint dim sv sh = do
  rss <- QC.generate $ mapM (const $ genRandomZeroOne) [1..]
  let (bs, w) = play rss dim sv sh
  putStr $ (unlines $ map (\(b,p) -> prettyShowGame (Game b p)) (zip bs (cycle [H,V]))) ++ "\nWinner: " ++ show w ++ " with " ++ show (lead (Game (last bs) w)) ++ "\n"