module Data.DP.Examples.Bigram where 
import Data.DP
import Data.DP.Solvers.TopDown
import Data.DP.SolverAPI
import Data.Semiring.Viterbi
import Data.Semiring.Derivation
import Data.Semiring.Prob
import Data.Semiring.ViterbiNBestDerivation
import Control.Monad.Identity
import qualified Data.Map as M
type Bigram = (String, String)

bigrams = 
   [(("a", "b"), 0.5), 
    (("b", "b"), 0.4), 
    (("b", "a"), 0.3), 
    (("c", "d"), 0.2), 
    (("d", "c"), 0.1),
    (("SOS", "a"), 0.2) 
   ] 

bigramsStartWith word =
    filter ((== word) . fst. fst ) bigrams 

ngram :: DP Int (M.Map String (ViterbiDerivation Prob [Bigram]))
ngram 0 = mkCell $ [mkItem "SOS" one]
ngram i = 
    getCell (i-1) (\(word, lastScore) -> 
        mkCell $ do
          (bigram,score) <- bigramsStartWith word 
          return $ mkItem (snd bigram) $ 
                 (lastScore `times` (constant $ mkViterbi $ Weighted (Prob score, mkDerivation [bigram]))))

runNgram = getResult $ runIdentity $ solveDP topDownMap 5 ngram