{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-}

-----------------------------------------------------------------------------
-- Copyright 2019, Advise-Me project team. This file is distributed under 
-- the terms of the Apache License 2.0. For more information, see the files
-- "LICENSE.txt" and "NOTICE.txt", which are included in the distribution.
-----------------------------------------------------------------------------
-- |
-- Maintainer  :  bastiaan.heeren@ou.nl
-- Stability   :  provisional
-- Portability :  portable (depends on ghc)
--
-- This module describes the underlying mechanism for generating evidence given constraints over attributes.
--
-----------------------------------------------------------------------------

module Recognize.Model.EvidenceBuilder
   ( EvBuilder
   , putAttributes, updateVar, getValueOf
   , setValueOf, storeEvidenceMaybe, buildEvidence, newVar
   , getAttributesPerStep, getVariables, putVariables
   ) where

import Control.Monad.State
import Data.Semigroup
import qualified Data.Map as M
import Domain.Math.Expr
import Recognize.Data.Attribute
import Bayes.Evidence
import Recognize.Expr.Symbols

-- | Builder that consumes attributes and produces evidence.
-- It Maintains a quadruple state:
--
--    1. Evidence               -- The generated evidence
--    2. [[Attribute]]          -- Remaining Attributes per step
--    3. M.Map String Expr      -- Maps variable names to expressions
newtype EvBuilder a = EvBuilder (State EvState a)

data EvState = EvState
   { evidence    :: Evidence
   , attributes  :: [[Attribute]]
   , variables   :: M.Map String Expr
   }

instance Functor EvBuilder where
  fmap f m = pure f <*> m

instance Applicative EvBuilder where
  pure = return
  (<*>) = ap

instance Semigroup (EvBuilder a) where
  EvBuilder ev1 <> EvBuilder ev2 = EvBuilder $ do
    st <- get
    if null $ fromEvidence $ evidence $ execState ev1 st
      then ev2
      else ev1

instance Monad EvBuilder where
  return = EvBuilder . return
  EvBuilder m >>= f = EvBuilder $ do
    a <- m
    let EvBuilder x = f a
    x

instance MonadState EvState EvBuilder where
  get = EvBuilder get
  put s = EvBuilder $ put s

-- | Creates a new wildcard variable
newVar :: String -> EvBuilder Expr
newVar s = do
  let mn = namedWildcard s
  updateVar s mn
  return mn

-- | Give an existing variable a new value
updateVar :: String -> Expr -> EvBuilder ()
updateVar s v = EvBuilder $ modify $ \st ->
   st { variables = M.insert s v $ variables st }

-- | Retrieve the value belonging to a variable
getValueOf :: Expr -> EvBuilder Expr
getValueOf e@(Sym s [Var n]) = EvBuilder $ do
  m <- gets variables
  case (isNamedWildcardSymbol s,M.lookup n m) of
    (True,Just e') -> return e'
    _ -> return e
getValueOf e = return e

-- | Expr wrapper over `updateVar`
setValueOf :: Expr -> Expr -> EvBuilder ()
setValueOf (Sym _ [Var n]) v = updateVar n v
setValueOf _ _ = return ()

buildEvidence :: EvBuilder () -> [[Attribute]] -> Evidence
buildEvidence (EvBuilder m) attrs =
   evidence $ execState m $ EvState mempty attrs M.empty

-- | Adds evidence to the state
storeEvidence :: Evidence -> EvBuilder ()
storeEvidence ev = EvBuilder $ modify $ \st ->
   st { evidence = evidence st <> ev }

storeEvidenceMaybe :: Maybe Evidence -> EvBuilder ()
storeEvidenceMaybe = maybe (return ()) storeEvidence

-- | Replaces the attributes in the state
putAttributes :: [[Attribute]] -> EvBuilder ()
putAttributes attrs = EvBuilder $ modify $ \st -> st { attributes = attrs }

getAttributesPerStep :: EvBuilder [[Attribute]]
getAttributesPerStep = EvBuilder $ gets attributes

getVariables :: EvBuilder (M.Map String Expr)
getVariables = EvBuilder $ gets variables

putVariables :: M.Map String Expr -> EvBuilder ()
putVariables vs = EvBuilder $ modify $ \st -> st { variables = vs }