-- Haskell
-- Author: Dan Hathaway
-- Brute-Forcing a winning strategy for the game
-- "Shut the box" using two rows of numbers:
--
-- 9 8 7 6 5 4 3 2 1
-- 1 2 3 4 5 6 7 8 9

----------------------------------------------------
--                   Importing
----------------------------------------------------

import Data.Map (Map)
import List
import Data.Maybe
import qualified Data.Map as Map
import IO

----------------------------------------------------
--                   Data types
----------------------------------------------------

-- Every entry in the board is either 0, 1, or 2
-- which represents the number of tiles in that
-- location that are up.
type Board = [Int]

-- A number associated with each board: rank = sum board
type Rank = Int

-- The sum of values on two dice.
type Roll = Int

data NumDice = OneDie | TwoDie
    deriving (Show,Eq)

-- "Board Values".
-- The "expected value" of each board.
-- This is the key quantity that is being computed.
type BV = Map Board (Float,NumDice)

----------------------------------------------------
--                  BASIC FUNCTIONS
----------------------------------------------------

-- The probability of rolling a number (sum of two dice):
pd :: Int -> Float
pd 0  = 0 -- Trivial cases.
pd 1  = 0 
pd 2  = 1/36
pd 3  = 2/36
pd 4  = 3/36
pd 5  = 4/36
pd 6  = 5/36
pd 7  = 6/36
pd 8  = 5/36
pd 9  = 4/36
pd 10 = 3/36
pd 11 = 2/36
pd 12 = 1/36

full_board :: Board
full_board = [2,2,2,2,2,2,2,2,2]

empty_board :: Board
empty_board = [0,0,0,0,0,0,0,0,0]

rank :: Board -> Int
rank b = sum b

----------------------------------------------------
--              COMPUTING THE PAYOFF
----------------------------------------------------

-- The score of a final state of the game.
-- (Called when "next_board board" is empty).
payoff :: Board -> Float
payoff b =
--    payoff_type2 b 0.0
    payoff_type1 b

-- The usual payoff (the score is the weighted
-- sum of the tiles that remain on the board).
payoff_type1 :: Board -> Float
payoff_type1 b =
    payoff1 (0,b)

-- The payoff in the type of game where the
-- goal is to have the tiles sum to less than
-- some number.
payoff_type2 :: Board -> Float -> Float
payoff_type2 b v =
    case payoff1 (0,b) <= v of
        True  -> 0
        False -> 1

-- A helper function for "payoff".
payoff1 :: (Int,Board) -> Float
payoff1 (i,b) =
    case b of
        v:t -> (payoff2 i v) + payoff1 (i+1,t)
        []  -> 0

-- A helper function for "payoff1".
-- i is 0, 1, ... or 8.
-- v is 0, 1, or 2.
payoff2 :: Int -> Int -> Float
payoff2 i v =
    case v of
        0 ->  0.0
        1 ->  9.0-i_f
        2 -> (9.0-i_f) + 2.0*(1.0+i_f)
    where
        i_f = fromIntegral i

----------------------------------------------------
--           COMPUTING THE NEXT BOARDS
----------------------------------------------------

-- Returns all boards that can be obtained
-- (in one move) from the given board with
-- a particular roll.
next_boards :: Board -> Int -> [Board]
next_boards b r =
    List.nub $ next_boards1 b r

next_boards1 :: Board -> Int -> [Board]
next_boards1 b 0 = [b]
next_boards1 b r =
    concat [try_move b n (r-n) | n <- [1..m]]
    where m = min r 9

-- Returns all boards that can be
-- obtained by first trying to shut
-- the tile with number n.
-- "rem_n" is the number that remains
-- to be handled after the appropriate
-- tile is shut.
try_move :: Board -> Int -> Int -> [Board]
try_move b n rem_n =
    (try_move_row1 b n rem_n) ++
    (try_move_row2 b n rem_n)

try_move_row1 :: Board -> Int -> Int -> [Board]
try_move_row1 b n rem_n =
    case (b !! pos) of
        2 -> next_boards nb rem_n
        _ -> []
    where
        pos = n-1
        nb = change_board b pos 1

try_move_row2 :: Board -> Int -> Int -> [Board]
try_move_row2 b n rem_n =
    case (b !! pos) of
        1 -> next_boards nb rem_n
        _ -> []
    where
        pos = 9-n
        nb = change_board b pos 0

change_board :: Board -> Int -> Int -> Board
change_board b pos v =
    a ++ (v:bs)
    where (a,_:bs) = splitAt pos b

----------------------------------------------------
--       ENUMERATING BOARDS OF A GIVEN RANK
----------------------------------------------------

boards_of_rank :: Int -> [Board]
boards_of_rank n =
    let cs = c_n_k 18 n
    in  nub (map boards_of_rank1 cs)
    
-- Helper function of boards_of_rank.
boards_of_rank1 :: [Int] -> Board
boards_of_rank1 x =
    let f i = quot i 2
    in  boards_of_rank2 (map f x)

-- Helper function of boards_of_rank1.
boards_of_rank2 :: [Int] -> Board
boards_of_rank2 [] = empty_board
boards_of_rank2 (x:xs) =
    let rec = boards_of_rank2 xs
    in  change_board rec x ((rec !! x)+1)

-- n choose k.
c_n_k :: Int -> Int -> [[Int]]
c_n_k n k = c_n_k_a n k 0

-- helper function for n choose k.
c_n_k_a :: Int -> Int -> Int -> [[Int]]
c_n_k_a n 0 a = [[]]
c_n_k_a 0 k a = []
c_n_k_a n k a =
    let h1 i = [(a+i):r | r <-
         (c_n_k_a (n-i-1) (k-1) (a+i+1))]
        res = concat [h1 i | i <- [0..(n-1)]]
    in  case k > n of
        True  -> []
        False -> res

----------------------------------------------------
--   COMPUTING THE EXPECTED VALUE OF EVERY BOARD
----------------------------------------------------

-- The main function.
compute_game :: BV
compute_game =
    compute_game1 18

-- Helper function of compute_game.
compute_game1 :: Int -> BV
compute_game1 0 =
    compute_ev_of_rank 0 Map.empty
compute_game1 n =
    let rec = compute_game1 (n-1)
    in  compute_ev_of_rank n rec

-- The driver function.
-- The input for this function is the rank
-- and the expected values of all
-- boards of rank < n.
compute_ev_of_rank :: Int -> BV -> BV
compute_ev_of_rank n bv =
    let bs = boards_of_rank n
        bv1 = [(b, compute_ev bv b) | b <- bs]
    in  Map.union bv (Map.fromList bv1)

-- The core function.
compute_ev :: BV -> Board -> (Float,NumDice)
compute_ev bv b =
    let en = exposed_num b
        d_ev  = compute_double_ev bv b
        s_ev1 = compute_single_ev bv b
        s_ev = if en <= 6 then s_ev1 else 200.0
    in  if d_ev < s_ev then (d_ev,TwoDie) else (s_ev,OneDie)

-- Computes the expected value given
-- that one dice is rolled.
compute_single_ev :: BV -> Board -> Float
compute_single_ev bv b =
    sum [(1/6) * compute_ev1 bv b r | r <- [1..6]]

-- Computes the expected value given
-- that two dice are rolled.
compute_double_ev :: BV -> Board -> Float
compute_double_ev bv b =
    sum [pd(r) * compute_ev1 bv b r | r <- [2..12]]

-- Helper function of compute_ev.
compute_ev1 :: BV -> Board -> Roll -> Float
compute_ev1 bv b r =
    let nb = next_boards b r
        f nbi = fst $ fromJust $ Map.lookup nbi bv
    in  case null nb of
        True  -> payoff b
        False -> minimum [f nbi | nbi <- nb]

-- The sum of the tiles that are showing.
exposed_num :: Board -> Int
exposed_num b =
    let f i v = case v of
            0 -> 0
            1 -> 9-i
            2 -> 1+i
    in  sum [f i (b !! i) | i <- [0..8]]

diff_boards :: Board -> Board -> [Int]
diff_boards b1 b2 =
    let f (x,y) = x-y
    in  map f (zip b1 b2)

----------------------------------------------------
--         COMPUTING THE BEST MOVE GIVEN
--         THE COMPUTED EXPECTED VALUES
----------------------------------------------------

-- NOTE: THE "Map.findMin" function is NOT APPROPRIATE!!!
best_move :: BV -> Board -> Roll -> ([Int],Float,Maybe NumDice)
best_move g b r =
    let nb = next_boards b r
        vs1 = [(b',fromJust $ Map.lookup b' g) | b' <- nb]
        (best_b,(ev,num_dice)) = best_move1 vs1 (head vs1)
    in  case null vs1 of
        True  -> (empty_board,payoff b,Nothing)
        False -> (diff_boards b best_b,ev,Just num_dice)

-- A dinky helper function for best_move.
best_move1 ::
    [(Board,(Float,NumDice))] -> 
    (Board,(Float,NumDice)) ->
    (Board,(Float,NumDice))
best_move1 [] best = best
best_move1 ((b,(f,nd)):xs) best =
    let new_best =
         if f < fst (snd best)
         then (b,(f,nd))
         else best
    in  best_move1 xs new_best

----------------------------------------------------
--         PRINTING COMPUTED EXPECTED VALUES
----------------------------------------------------

print_payload :: [Char]
print_payload =
    let game = compute_game1 18
        payload = Map.toList game
    in  concat [print_helper l | l <- payload]

-- Printer helper:
print_helper :: (Board,(Float,NumDice)) -> [Char]
print_helper (b,(f,nd)) =
    let pnd x = case x of
            OneDie -> "1D"
            TwoDie -> "2D"
    in  concat [(show i) ++ " " | i <- b]
        ++ (show f) ++ " "
        ++ (pnd nd) ++ "\n"

----------------------------------------------------
--         READING PRE-COMPUTED EXPECTED VALUES
----------------------------------------------------

-- Reads the contents of a file into
-- the "Board Values" map.
read_bv_from_file :: String -> BV
read_bv_from_file cts =
    foldl read_single_bv Map.empty (lines cts)

-- A helper function for reading in the
-- "Boad Values" Map from a file:
read_single_bv :: BV -> String -> BV
read_single_bv bv line =
    let ss = words line
        b = [ (read x)::Int | x <- (take 9 ss) ]
        e = (read (ss !! 9))::Float
        nd = case ss !! 10 of
            "1D" -> OneDie
            "2D" -> TwoDie
    in  Map.insert b (e,nd) bv

----------------------------------------------------
--         HOW TO USE THE PROGRAM PART 1
----------------------------------------------------
-- How to print the winning strategy to a file:

-- 1) Uncomment the line below.
-- 2) Be in the same CWD as this file.
-- 3) Run the command "ghc --make ShutTheBox.hs".
-- 4) Run "ShutTheBox.exe > Data/game.txt".

main = putStr print_payload

----------------------------------------------------
--         HOW TO USE THE PROGRAM PART 2
----------------------------------------------------
-- How to find the optimal move given the
-- precomputed board expected values file:

-- 1)  Be in the same CWD as this file.
-- 2)  Make sure this no line "main = ..." in this file.
-- 3)  Run the command "ghci"
-- 4)  Type: :l ShutTheBox
-- 5)  Type: game_file <- readFile "Data/full_game.txt"
-- 6)  Type: let game = read_bv_from_file game_file
-- 7)  (for fun) Type: Map.lookup full_board game
--     It may take a few seconds for this command to
--     run because haskell is lazy.
--     The first value is the expected score of the game
--     (assuming you play optimally) and the second value
--     tells you how many dice you should roll if that is
--     the current board.
-- 8)  Start playing the game "Shut The Box" with a friend
--     and roll two 6 sided die for the first turn.
--     Let r denote their sum.
-- 9)  Type: best_move game full_board r
--     This will tell you which tiles to put down and how
--     many dice to roll the next turn.
--     Do what the function tells you to do.
-- 10) Keeping rolling the die repeating step 9
--     (using the current board as input)
--     until the game ends.
-- 11) Keep playing the game over and over again until
--     your friend gets angry that your total score
--     (carried over from all the games you play) is
--     lower than his.
-- 12) Inform your friend that he has lost "The Game",
--     which you just lost too if you are are nerd
--     and are reading this.