-- | Checkerboard problem from <http://en.wikipedia.org/wiki/Dynamic_programming>
--   This example gives an implementation that matches the wikipedia example in 
--   speed, with a third as much code and much more generality. 

module Data.DP.Examples.CheckerBoard where 

import Data.DP
import Data.DP.Solvers.TopDown
import Data.DP.Solvers.Recursive
import Data.DP.SolverAPI
import Data.Semiring.Max
import Data.Semiring.ViterbiNBestDerivation
import Data.Semiring.Derivation
import Data.Semiring.Viterbi
import Control.Monad.Identity

checkerScore :: [[Int]]
checkerScore = 
    reverse $  [[6, 7, 4, 7, 8],
                [7, 6, 1, 1, 4],
                [3, 5, 7, 8, 2],
                [0, 6, 7, 0, 0],
                [0, 0, 5, 0, 0]]
    
getScore (i,j)  = (checkerScore !! (i-1)) !! (j-1) 

n = 5

data CheckerState = Finish | Middle (Int, Int)
                  deriving (Eq, Ord)


checkerBoard ind = 
    case ind of 
      Finish -> mconcat $ map (\j-> f' (n,j)) [1..n]
      (Middle (i,j)) -> 
          if j < 1 || j > n then mempty
          else if i == 1 then constant $ Max $ getScore (i,j)
          else  (mconcat $ [f' (i-1,j-1), f' (i-1,j), f' (i-1,j+1)]) 
              `times` (constant $ Max $ getScore (i,j))
    where f' = f . Middle

runCheckerboard = getSimpleResult $ runIdentity $ solveSimpleDP topDownMap Finish checkerBoard 

checkerBoardGen mkSemi ind = 
    case ind of 
      Finish -> mconcat $ map (\j-> f' (n,j)) [1..n]
      (Middle (i,j)) -> 
          if j < 1 || j > n then mempty
          else if i == 1 then constant $ mkSemi (i,j)
          else  (mconcat $ [f' (i-1,j-1), f' (i-1,j), f' (i-1,j+1)]) 
              `times` (constant $ mkSemi (i,j))
    where f' = f . Middle

maxSemi :: (Int, Int) -> Max Int
maxSemi = Max . getScore

maxSemiViterbi :: (Int, Int) -> ViterbiDerivation (Max Int)  [(Int,Int)]
maxSemiViterbi pos = mkViterbi $ Weighted (Max $ getScore pos, mkDerivation [pos])

runCheckerboardGen semi = getSimpleResult $ runIdentity $ solveSimpleDP recursive Finish (checkerBoardGen semi)