-----------------------------------------------------------------------------
-- 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 Task.TheatreRate.Recognizer (theatreRate) where

import Control.Arrow
import Control.Applicative                      (many,empty)
import Control.Monad                            (guard, msum)
import Data.Maybe
import Util.List
import Util.Monad
import Recognize.Data.Approach
import Recognize.Data.Math
import Recognize.Data.MathParserOptions
import Recognize.Data.MathParserOutput
import Recognize.Data.MathStoryProblem
import Recognize.Data.Diagnosis
import Recognize.Data.StringLexer
import Recognize.Expr.Functions
import Recognize.Parsing.Derived
import Recognize.Parsing.Parse
import Recognize.Recognizer
import Recognize.Strategy.Applications
import Recognize.SubExpr.Recognizer
import Recognize.SubExpr.SEParser
import Recognize.SubExpr.Symbols
import Domain.Math.Data.Relation
import Domain.Math.Expr.Data
import Ideas.Common.Id
import Ideas.Utils.Uniplate                     (transform)
import Ideas.Utils.Prelude
import Task.TheatreRate.Assess
import Task.Network.TheatreRate
import Bayes.Evidence ( evidenceOfAbsence )

theatreRate :: MathStoryProblem
theatreRate = mathStoryProblem
   { problemId     = newId "theatrerate"
   , processInputs = filter ((== newId "07b") . getId)
   , analyzers     = [(newId "07b", ana)]
   , inputFile     = Just "input/theatrerate.csv"
   , networkFile   = Just "networks/TheatreRate.xdsl"
   , singleNetwork = network
   }
 where
   ana = analyzer
      { parser     = mathParser mathParserOptions {multByConcatenation = False} . stringLexerOutput
      , recognizer = seRecognizer pDiagnosis . mathParserOutput . modifyInput
      , collector  = evidenceOfAbsence ans2 False . assess'
      }

-- Somewhat of a hack, since the lexer/parser should correctly handle whitelisting of variables
-- Should be fixed once pilots are finished
modifyInput :: MathParserOutput -> MathParserOutput
modifyInput (MathParserOutput mpo che) = MathParserOutput (map math mpo) che
  where
    math (M t ethe) = M t $ fmap (transform exprDiv . transform expr) ethe
    expr e = case e of
      (Var "r" :*: 1) -> Var "r1"
      (Var "r" :*: 2) -> Var "r2"
      (Var "R" :*: 1) -> Var "R1"
      (Var "R" :*: 2) -> Var "R2"
      (Var "T" :*: 1) -> Var "T1"
      (Var "T" :*: 2) -> Var "T2"
      (Var "t" :*: 1) -> Var "t1"
      (Var "t" :*: 2) -> Var "t2"
      _ -> e
    -- Also a quickfix for variables ending with a colon being recognized as division
    exprDiv e = case e of
      (Var "R1" :/: e') -> e'
      (Var "r1" :/: e') -> e'
      (Var "R2" :/: e') -> e'
      (Var "r2" :/: e') -> e'
      (Var "T1" :/: e') -> e'
      (Var "t1" :/: e') -> e'
      (Var "T2" :/: e') -> e'
      (Var "t2" :/: e') -> e'
      _ -> e

pDiagnosis :: SEParser Diagnosis
pDiagnosis = do
  (appr, st) <- pSteps
  let   sd = newDiagnosis appr st
  --    exprs = rights (map (getResult . getMath) st)
  --    exprAsString = concatMap show exprs
  return sd


pSteps :: SEParser (Approach, [Step])
pSteps = do
  (ap,e,st,sk) <- choice
    [
      (\(e,st,sk) -> (Numerical,e,st,sk)) <$> withGuard (not . null . snd3) pStepsNumerical
    , (\(e,st,sk) -> (Algebraic,e,st,sk)) <$> withGuard (not . null . snd3) pStepsAlgebraic
    ]
  fa_st <- pFinalAnswer sk e
  let ss = st ++ maybeToList fa_st
  guard (length ss >= 1)
  return (ap,ss)

r1_def :: Expr -> Expr
r1_def s = 4*s + 30 <!> 30*s + 4 <!> 30*s

r2_def :: Expr -> Expr
r2_def s = 8*s <!> newMagicNat*s

r12_def :: Expr
r12_def = lt "s" newMagicVar $ \s -> lbl "R1" (r1_def s) <?> lbl "R2" (r2_def s)

r12_equation :: Expr
r12_equation = lbl "setup equation" $ lt "s" newMagicVar $ \s -> r1_def s <&> r2_def s

r12_inequation :: Expr
r12_inequation = lbl "setup inequation" $ lt "s" newMagicVar $ \s -> r1_def s <&> r2_def s

-- | Linearly solve the answer
pStepsAlgebraic :: SEParser (Expr,[Step],[Math])
pStepsAlgebraic =
  pInOrder
    [ \_ -> do
      -- T1: 30+4x
      -- T2: 8x
      modify $ \st -> st { inputType = Just [Expr, Definition] }
      pMatchSubSteps r12_def
    , \mres -> case mres of
        Nothing -> empty
        Just (e1,st1) -> do
          (e2,st2) <- resetAfter $ pMatchSubSteps r12_def
          -- Make sure that we didn't match the same formula twice
          guard (e1 /= e2)
          -- Returning e2 only because we have to return an Expr
          return (e2,st1++st2)
    , \_ -> do
      -- 30+4x = 8x
      modify $ \st -> st { inputType = Just [LinearWithType EqualTo] }
      res <- resetAfter (pMatchSubSteps r12_equation)
      pLog ("test: " ++ show res)
      return res
    , \_ -> do
      -- 30+4x < 8x
      modify $ \st -> st { inputType = Just [LinearWithType LessThan, LinearWithType GreaterThan] }
      res <- resetAfter (pMatchSubSteps r12_inequation)
      pLog ("test: " ++ show res)
      return res
    ,
     \mres -> do
      pLog ("equation succeeded: " ++ show mres)
      mrel <- case mres of
                  Just (e,st) -> do
                    (e1,e2) <- getBinary e
                    let rel = e1 .==. e2
                    math <- maybeToParse $ safeLast $ getMaths st
                    return $ Just (rel,math)
                  Nothing -> return Nothing
      pLog $ "found rel: " ++ show mrel
      res <- pSolveLinear mrel
      guard (length (snd res) >= 2)
      return $ first rightHandSide res
    ]

-- | Fill in some values in the equations and see where it gets you...
pStepsNumerical :: SEParser (Expr, [Step],[Math])
pStepsNumerical = do
  pLog "Numerical"
  res <- many1 (pSkipUntil $ pMatchSubSteps nexpr)
  let sts = concatMap (snd . fst) res
  let sk = concatMap snd res
  let e = fst $ fst $ head res
  return (e, sts, sk)
  where
    nexpr = lt "r" newMagicNat $ \r -> lblE "Try" r ((30 + 4*r <!> 30*r + 4 <!> 30*r) <&> 8*r)

pFinalAnswer :: [Math] -> Expr -> SEParser (Maybe Step)
pFinalAnswer skipped e = do
  pLog "pFinalAnswer"
  rest <- many skip
  let answers = mapMaybe mAnswer (skipped ++ rest)
  let me = closestInList (filter (\n -> isNat n && n >= 6 && n <= 10) $ e : answers) 8
  case me of
    Nothing -> pLog "empty" >> return Nothing
    Just fe -> return $ Just $ makeFAStep fe

mAnswer :: Math -> Maybe Expr
mAnswer m =
  msum
    [ do
      (x :==: y) <- getEq m
      guard (isVar x && isNat y)
      return y
    , do
      n <- getExpr m
      guard (isNat n)
      return n
    ]