-----------------------------------------------------------------------------
-- 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)
--
-----------------------------------------------------------------------------

module Bayes.Script
   ( taskConcepts
   , emptyConcepts
   , aggregateConcepts
   , aggregateConcepts'
   , findEvidence
   , addModelToLatestRequest
   ) where

import Bayes.Evidence
import Bayes.EliminationOrdering ( minFactorOrder )
import Bayes.Inference ( infer, toEvidence )
import Bayes.Network
import Data.List hiding (insert)
import Data.Maybe
import Data.Function ( on )
import Control.Monad ( (>=>) )
import Database.HDBC
import Database.HDBC.Sqlite3
import Ideas.Common.Id ( showId )
import qualified Ideas.Text.XML as XML ( fromXML, XML, parseXML, findChild )
--import qualified Ideas.Text.XML.Interface as XML
import qualified Data.Map as M
import qualified Data.Set as S

import qualified Task.Network.StudentModel as StudentModel
import Recognize.Data.MathStoryProblem (Task(Task), singleNetwork)
import Bayes.Probability
import Database.Data
import Main.Tasks ( tasks, findTaskFuzzy )


-- | Add information from the student modeller to the database, based on the
-- raw input and output of the domain reasoner. This is unfortunately done in a
-- roundabout way, due to the way the framework functions.
--
-- We assume that the latest row to the requests table is the one we're dealing
-- with, since we must know the row ID to continue.
addModelToLatestRequest :: Connection -> IO ()
addModelToLatestRequest conn = do
   r <- fromMaybe (error "No request to add model to.") <$> latestRecord conn
   let xmlInput = either (error "Could not read XML") id . XML.parseXML $ r ! "input"
       (sID, tID) = fromMaybe mempty . collectIDs $ xmlInput
       inputIDs = intercalate "," . map fst . fromMaybe mempty . collectInputs $ xmlInput
       reqnr = r ! "rowid" :: Int
   ev <- findEvidence $ r ! "output"
   let partialModel = taskConcepts tID ev
   partialModels <- latestPartialModels conn sID
   let studentModel = aggregateConcepts $ (tID, partialModel) : partialModels
   insertRecord conn studentsTable
       ( insert "studentid" sID
       . insert "taskid" tID
       . insert "inputs" inputIDs
       . insert "requestnr" reqnr
       . insert "evidence" ev
       . insert "studentmodel" studentModel
       . insert "partialmodel" partialModel
       $ mempty )
   commit conn

-- Find evidence from a reply of the domain reasoner.
findEvidence :: Monad m => XML.XML -> m Evidence
findEvidence = XML.findChild "diagnosis" >=> XML.findChild "evidence" >=> XML.fromXML


-- | Run the evidence on the Bayesian networks to add soft evidence.
--
-- Instead of running it on the combined model, we run it through the models
-- seperately and only output the evidence as it should be in the student model
-- nodes. This should work since there should be no feedback effects at that
-- point.
--feed :: Evidence -> StudentModel
--feed = toEvidence StudentModel.network . feed'


-- | Obtain task concept node probabilities for a particular task.
taskConcepts :: TaskID -> Evidence -> Evidence
taskConcepts tID ev = 
   let nw = maybe (error $ "no such task " ++ tID) (\(Task t) -> singleNetwork t) $ findTaskFuzzy tID
   in  filterEvidence isConcept . toEvidence nw $ infer nw ev minFactorOrder (S.filter isConcept (nodeIds nw))


-- | Feeding empty evidence to the task concept evidence. This is used in case
-- the task is not even mentioned in the precalculations (presumably because no
-- request for the task was sent).
emptyConcepts :: [ (TaskID, Evidence) ]
emptyConcepts = flip map tasks $ \t -> let tID = showId t in (tID, taskConcepts tID mempty)

-- | Same as 'aggregateConcepts', but do not include task concept node
-- probabilities in final evidence.
aggregateConcepts' :: [ (TaskID, Evidence) ] -> Evidence
aggregateConcepts' precalculatedTaskConcepts 
   = toEvidence StudentModel.network
   . M.map ((\x->[x,1-x]) . bayesAverage) 
   . M.fromListWith (++) $
      [ (concept, [fromDouble . fromRational $ probability])
      | Task t <- tasks
      , let taskConcepts' = fromJust . lookup (showId t) $ precalculatedTaskConcepts ++ emptyConcepts
      , (taskConcept, stateLabel, probability) <- allProbabilities taskConcepts'
      , stateLabel == "Yes" || stateLabel == "State0" || stateLabel == "Correct" || stateLabel == "Correctr"
      , concept <- maybeToList . getConcept $ taskConcept
      ]


-- | Aggregate task concept node probabilities into generic concept node
-- probabilities.
aggregateConcepts :: [ (TaskID, Evidence) ] -> Evidence
aggregateConcepts p = aggregateConcepts' p `mappend` mconcat (map snd . nubBy ((==) `on` fst) $ p ++ emptyConcepts)


-- for now, just use plain average
bayesAverage :: [Probability] -> Probability
bayesAverage ps = sum ps / fromIntegral (length ps)

-- | Task model concept node ID to a student model concept node ID.
getConcept :: String -> Maybe String
getConcept ct = 
   let cs = takeWhile (/= '_') $ ct
   in if cs `S.member` nodeIds StudentModel.network
         then Just cs
         else Nothing

isConcept :: String -> Bool
isConcept = isJust . getConcept