{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- 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 defines some common functions used in the exercise specific assessment functions.
--
-----------------------------------------------------------------------------

module Recognize.Model.Assess
   ( -- * Generate evidence
     generateEvidence, answerCorrect, stringNode
     -- * Failure constraints
   , failOnInvalidEquationAny, failOnInvalidEquation, failOnIncorrectDistribution, failOnCommonMistake
   , failOnUnequalRelationAny, failOnUnequalRelation, failOnAtomMixedUpAny, failOnAnyMistake
   , withoutFailure
     -- * Evidence collecting
   , giveNodeAndCollect, giveNodeAndCollectAll, giveNodeAndCollectAllKnown, giveNodeAndCollectDefault
   ) where

import Control.Monad.Identity
import Control.Monad.State
import Data.Either
import Data.Maybe
import Data.Monoid
import Data.List
import Domain.Math.Expr
import Domain.Math.Data.Relation
import Ideas.Utils.Uniplate
import Recognize.Model.Constraint
import Recognize.Model.EvidenceBuilder
import Recognize.Model.Result
import Bayes.Evidence
import Recognize.Data.Approach
import Recognize.Data.Attribute
import Recognize.Model.Connectives
import Recognize.Expr.Normalform
import Recognize.Expr.Symbols
import Recognize.Expr.Functions
import Bayes.Network
import Util.List

-- | Generic function for gathering evidence
generateEvidence :: (Approach -> EvBuilder ()) -- ^ Function for building the evidence
              -> Approach -- ^ Category of solution
              -> [[Attribute]] -- ^ Collected attributes to use as input
              -> Evidence
generateEvidence feb appr = buildEvidence (feb appr)

--Produce evidence for the correct answer from the attributes, with the correct answer being the given expression
answerCorrect :: Expr -> [[Attribute]] -> Node Bool -> Evidence
answerCorrect ans attrs aC = mconcat (catMaybes [ verifyConstraintI (answerConstraint ans aC) (concat attrs)])

--Produce evidence for a string Node based on a string:
stringNode :: Maybe String -> Node String -> Evidence
stringNode (Just s) n = n .== s
stringNode Nothing _  = mempty

--Constraint that checks if a finalanswer matches the given expression
answerConstraint :: Expr -> Node Bool -> NodeConstraint Bool Identity [Attribute]
answerConstraint fa aC =
  makeNodeConstraint aC $ \xs -> return $
    case xs of
      _ | FinalAnswer fa `elem` xs -> Success
        | hasFA xs -> Failure
        | otherwise -> Unknown
    where
      hasFA xs = isJust $ find (\attr -> case attr of
                        FinalAnswer _ -> True
                        _ -> False) xs

-- | Generates a failure if some InvalidEquation attribute is present
failOnInvalidEquationAny :: Constraint EvBuilder [Attribute]
failOnInvalidEquationAny = negateConstraint (exists1 $ InvalidEquation wildcard wildcard)

-- | Generates a failure if an InvalidEquation attribute with the two specified expression is present
failOnInvalidEquation :: Expr -> Expr -> Constraint EvBuilder [Attribute]
failOnInvalidEquation e1 e2 = negateConstraint (exists1 $ InvalidEquation e1 e2)

-- | Generates a failure if the IncorrectDistribution attribute or the NonMatchingParenthese attribute is present
failOnIncorrectDistribution :: Constraint EvBuilder [Attribute]
failOnIncorrectDistribution = negateConstraint (exists1 IncorrectDistribution <||> exists1 NonMatchingParentheses)

-- | Generates a failure if the CommonMistake attribute is present
failOnCommonMistake :: Constraint EvBuilder [Attribute]
failOnCommonMistake = negateConstraint (exists1 CommonMistake)

-- | Generates a failure if the UnequalRelations attribute with the two specified relations is present
failOnUnequalRelation :: Relation Expr -> Relation Expr -> Constraint EvBuilder [Attribute]
failOnUnequalRelation rel1 rel2 = negateConstraint (exists1 $ UnequalRelations rel1 rel2)

-- | Generates a failure if the UnequalRelations attribute is present
failOnUnequalRelationAny :: Constraint EvBuilder [Attribute]
failOnUnequalRelationAny = failOnUnequalRelation wildcardRelation wildcardRelation

-- | Generates a failure if the AtomMixedUp attribute is present
failOnAtomMixedUpAny :: Constraint EvBuilder [Attribute]
failOnAtomMixedUpAny = negateConstraint (exists1 $ AtomMixedUp wildcard wildcard)

-- | Combines all the above mentioned failure functions
failOnAnyMistake :: Constraint EvBuilder [Attribute]
failOnAnyMistake = failOnCommonMistake <> failOnInvalidEquationAny <> failOnUnequalRelationAny <> failOnAtomMixedUpAny <> failOnIncorrectDistribution

-- | Combines the given constraint with the constraint that there should be no failures
withoutFailure :: Constraint EvBuilder [Attribute] -> Constraint EvBuilder [Attribute]
withoutFailure c = c ==> (failOnAnyMistake <?>> mempty)

-- | Collect/Store evidence generated by a node constraint
collect :: Verifiable b => NodeConstraint b EvBuilder [Attribute] -> EvBuilder ()
collect c = verifyConstraint c >>= storeEvidenceMaybe

-- | Collect/Store evidence generated by a node constraint
-- If the evidence is unknown then we assign a value to the given assessment variable
collectDefault :: Verifiable b => NodeConstraint b EvBuilder [Attribute] -> Expr -> Expr -> EvBuilder ()
collectDefault c var new = do
  mev <- verifyConstraint c
  val <- (nfComAssoc . normalizeIfNF) <$> transformM getValueOf new
  replaceWhenUnknown mev var val
  storeEvidenceMaybe mev

replaceWhenUnknown :: Maybe a -> Expr -> Expr -> EvBuilder ()
replaceWhenUnknown Nothing var new = setValueOf var new
replaceWhenUnknown _ _ _ = return ()

-- | Takes the given constraint, couples it to the node and stores the evidence.
giveNodeAndCollect :: Verifiable b => Node b -> Constraint EvBuilder [Attribute] -> EvBuilder ()
giveNodeAndCollect n c = do
  let nodeconstraint = giveNode n c
  collect nodeconstraint

-- | Useful when multiple constraints spread over multiple steps have to be verified.
-- Takes the given constraints and couples them to the node.
-- The evidence is true if ALL given results are successful. (= no unknwon results)
-- The evidence is false if ANY known result is false.
-- The evidence is unknown if all known results are succesful, but there are unknown results.
giveNodeAndCollectAll :: Verifiable b => Node b -> [Constraint EvBuilder [Attribute]] -> EvBuilder ()
giveNodeAndCollectAll n cs = do
  evs <- mapM (verifyConstraint . giveNode n) cs
  let knownEvs = lefts (map snd (concatMap fromEvidence (catMaybes evs)))
  let nodeconstraint | 1 `elem` knownEvs = giveNode n failure
                     | any isNothing evs = giveNode n unknown
                     | otherwise = giveNode n success
  storeEvidenceMaybe (verifyConstraintI nodeconstraint [])

-- | Useful when multiple constraints spread over multiple steps have to be verified.
-- Takes the given constraints and couples them to the node.
-- The evidence is true if ALL known results are successful.
-- The evidence is false if ANY known result is false.
-- The evidence is unknown if there are no known results.
giveNodeAndCollectAllKnown :: Verifiable b => Node b -> [Constraint EvBuilder [Attribute]] -> EvBuilder ()
giveNodeAndCollectAllKnown n cs = do
  evs <- mapM (verifyConstraint . giveNode n) cs
  let knownEvs = lefts (map snd (concatMap fromEvidence (catMaybes evs)))
  let nodeconstraint | null knownEvs = giveNode n unknown
                     | 1 `elem` knownEvs = giveNode n failure
                     | otherwise = giveNode n success
  storeEvidenceMaybe (verifyConstraintI nodeconstraint [])

-- | Takes the given constraint, couples it to the node and stores the evidence.
-- If the evidence is unknown then we assign a value to the given assessment variable
giveNodeAndCollectDefault :: Verifiable b => Node b -> Constraint EvBuilder [Attribute] -> Expr -> Expr -> EvBuilder ()
giveNodeAndCollectDefault n c var new = do
  let nodeconstraint = giveNode n c
  collectDefault nodeconstraint var new

-- | Type class for all types of nodes that we can verify using a constraint.
class Verifiable b where
    verifyConstraint :: NodeConstraint b EvBuilder [Attribute] -> EvBuilder (Maybe Evidence) -- ^ Verify a node constraint using the evidence builder.
    verifyConstraintI :: NodeConstraint b Identity [a] -> [a]-> Maybe Evidence -- ^ Verify a node constraint without consuming attributes.

instance Verifiable Bool where
    verifyConstraint c = do
           attrs <- getAttributesPerStep
           -- For each step that produced attributes we attempt to verify the constraint with
           mi <- returnOnJustM (\a -> do
              st <- get
              res <- getResult (getConstraint c) a
              when (isUnknown res) $ put st
              return $ resultToBool res) attrs
           case mi of
              Nothing -> do
                 r <- getResult (getConstraint c) []
                 return (resultToBool r >>= \i -> return $ getNode c .== i)
              Just (i,attrs') -> do
                 putAttributes attrs'
                 return $ Just $ getNode c .== i

    verifyConstraintI c attrs =
       case resultToBool $ runIdentity $ getResult (getConstraint c) attrs of
          Nothing -> Nothing
          Just i  -> Just $ getNode c .== i

instance Verifiable (Maybe Bool) where
    verifyConstraint c = do
           attrs <- getAttributesPerStep
           -- For each step that produced attributes we attempt to verify the constraint with
           mi <- returnOnJustM (\a -> do
              st <- get
              res <- getResult (getConstraint c) a
              when (isUnknown res) $ put st
              return $ resultToBool res) attrs
           case mi of
              Nothing -> do
                 r <- getResult (getConstraint c) []
                 return (resultToBool r >>= \i -> return $ getNode c .== (Just i))
              Just (i,attrs') -> do
                 putAttributes attrs'
                 return $ Just $ getNode c .== (Just i)

    verifyConstraintI c attrs =
       case resultToBool $ runIdentity $ getResult (getConstraint c) attrs of
          Nothing -> Nothing
          Just i  -> Just $ getNode c .== (Just i)