-----------------------------------------------------------------------------
-- 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.Pattern.Recognizer (pattern) where

import Control.Arrow
import Control.Applicative (empty,many)
import Control.Monad (guard, msum)
import Data.Maybe
import Util.List
import Util.Monad
import Recognize.Data.Attribute hiding (Other)
import Recognize.Data.Math
import Recognize.Data.MathStoryProblem
import Recognize.Data.MathParserOptions
import Recognize.Data.MathParserOutput
import Recognize.Data.Diagnosis
import Recognize.Data.Approach
import Recognize.Data.Step
import Recognize.Data.StringLexer
import Recognize.Expr.Functions
import Recognize.Expr.Normalform
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 hiding (ltSymbol)
import Task.Pattern.Assess
import Domain.Math.Data.Relation
import Domain.Math.Expr.Data
import Ideas.Common.Id
import Ideas.Common.Rewriting
import Ideas.Text.OpenMath.Dictionary.Relation1
import Task.Network.Pattern
import Bayes.Evidence ( evidenceOfAbsence )

pattern :: MathStoryProblem
pattern = mathStoryProblem
   { problemId   = newId "pattern"
   , analyzers   = [(newId "04", ana)]
   , inputFile   = Just "input/pattern.csv"
   , networkFile = Just "networks/Pattern.xdsl"
   , singleNetwork = network
   }
 where
   ana = analyzer
      { parser     = mathParser mathParserOptions { functionCallWhitelist = "uU" } . stringLexerOutput
      , recognizer = seRecognizer pDiagnosis . mathParserOutput . modifyInput
      , collector  = evidenceOfAbsence ans1 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 remVar ethe
    -- For this exercise we wish to remove expressions such as 15*th, which could have been written as 15th
    -- We are fairly certain that in this exercise the variables in singular expressions such as the one above
    -- are meaningless and may be removed. If this appears not the be the case, then this function should be removed.
    remVar (e :*: Var _) = e
    remVar 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 (\(_,xs,_) -> not $ null xs) pStepsNumerical
    , (\(e,st,sk) -> ((Other "Numerical2"),e,st,sk))    <$> withGuard (\(_,xs,_) -> not $ null xs) pStepsRecursive
    , (\(e,st,sk) -> (Generalizing,e,st,sk)) <$> pStepsGeneralizing
    ]
  fa_st <- pFinalAnswer sk e
  let ss = st ++ maybeToList fa_st
  guard (length ss >= 1)
  return (ap,ss)

-- | This is a counting strategy
--
-- Usually starts at 4,7 or 10 and then continuously adds 3 (unless a mistake is made)
pStepsNumerical :: SEParser (Expr, [Step],[Math])
pStepsNumerical = do
      pInOrderAll [ const $ withGuard (\(_, xs) -> length xs >= 3) $ pNumAndVals 4 ]

-- | Parses the correct values in combination with numbers
pNumAndVals :: Expr -> SEParser (Expr, [Step])
pNumAndVals e = do
  me <- option $ pNumAndVal e
  case me of
    Nothing -> return (e,[])
    Just (e',ss) -> second (ss:) <$> pNumAndVals e'
  where
    pNumAndVal ex = do
      (b',attr,math) <- choice
        [ do
          -- No numbers, so:
          -- 4
          -- 7
          -- 10
          -- ..
          m_a <- skip
          a <- getExpr m_a
          (b',attr) <- isVal ex a
          return (b', attr, m_a)
        , do
          -- 1 == 4
          -- 2 == 7
          -- 3 == 10
          -- ..
          meq <- skip
          (a :==: b) <- getEq meq
          guard (isNat a)
          (b',attr) <- isVal ex b
          return (b', attr, meq)
        , do
          -- 1 4
          -- 2 7
          -- 3 10
          -- ..
          m_a <- skip
          a <- getExpr m_a
          guard (isNat a)
          m_b <- skip
          b <- getExpr m_b
          (b',attr) <- isVal ex b
          return (b', attr, m_b)]
      return (b', valStep attr math)

isVal :: Expr -> Expr -> SEParser (Expr, [Attribute])
isVal e1 e2 =
  maybeToParse $ msum
    [ do
      let diff = nf $ (e2 - e1) / 3
      guard $ isNat diff && diff <= 5 && not (e1 == 4 && e2 == 16)
      return (e2,[NExpr e2])
    , do
      -- Added 4 instead of 3
      let diff = nf $ (e2 - e1) / 4
      guard $ isNat diff && diff <= 5 && not (e1 == 4 && e2 == 16)
      return (e2,[CommonMistake])
    , do
      -- Added 2 instead of 3
      let diff = nf $ (e2 - e1) / 2
      guard $ isNat diff && diff <= 5 && not (e1 == 4 && e2 == 16)
      return (e2,[CommonMistake])
    , do
      -- 21,24,27,54
      -- Sometimes halfway they multiply by two
      guard $ nf (e2 / e1) == 2
      return (e2,[CommonMistake])
    ]

valStep :: [Attribute] -> Math -> Step
valStep attr m = Step (newId "num") (m,attr) []

-- | Calculates the answer by dividing the number of tiles to go divided by the number of tiles that are added each pattern increase
pStepsRecursive :: SEParser (Expr, [Step],[Math])
pStepsRecursive = do
  n <- getInputSize
  pLog $ "Size: " ++ show n
  pInOrder
    [ \_ -> pMatchSubSteps rexpr2
    , \m -> case m of
        (Just (a,_)) -> pMatchSubSteps $ rexpr3_a a
        Nothing -> pMatchSubSteps rexpr3_b
    , \_ -> pMatchSubSteps rexpr4
    ]
    where
      rexpr2 = lbl "subtract" $ newMagicNat - newMagicNat
      rexpr3_a e = lbl "divide" $ e / (3 <!> 2 <!> 4)
      rexpr3_b = lbl "divide" $ sub $ ((46 <?> 43 <?> 40) <!> newMagicNat) / (3 <!> 2 <!> 4)
      rexpr4 = lbl "increase" $ lt "u" newMagicNat $ \u -> sub (u + sim (16 - u))

-- | Linearly solves the answer
pStepsGeneralizing :: SEParser (Expr, [Step],[Math])
pStepsGeneralizing = do
  pInOrder
    [ \_ -> pMatchSubSteps gstep1
    , \_ -> do
      -- Ensures only one input will be consumed
      modify $ \st -> st
         { optTraverse = False
         , optIterate  = False
         , inputType   = Just [Linear, Definition]
         }
      resetAfter (pMatchSubSteps gstep2_def)
    , \_ -> do
      -- 4+3x = 50
      let
      modify $ \st -> st { inputType = Just [LinearWithType EqualTo] }
      res <- resetAfter (pMatchSubSteps gstep2_eq)
      return res
    , \_ -> do
      -- 4+3x < 50
      modify $ \st -> st { inputType = Just [LinearWithType LessThan, LinearWithType GreaterThan] }
      res <- resetAfter (pMatchSubSteps gstep2_ineq)
      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
    , \res -> case res of
      Nothing -> empty
      Just (e,_) -> pMatchSubSteps (gstep3 e)
    ]
    where
      gstep1 = 7 <?> (lbl "add 3 tiles" 3) + 4
      gstep2 = ((4 <!> 7) + 3*newMagicVar) <!> newMagicNat + 3*newMagicVar <!>  7*newMagicVar
      gstep2_def = lbl "definition" gstep2
      gstep2_eq = lbl "equation" (gstep2 <&> (50 <!> newMagicNat))
      gstep2_ineq = lbl "inequation" (gstep2 <&> (50 <!> newMagicNat))
      gstep3 r = (lbl "pattern" $ (16 <?> ceilingExpr (nf4 2 r)) <!> floorExpr (nf4 2 r)) <?> (lbl "solution" $ nf4 2 r)

pFinalAnswer :: [Math] -> Expr -> SEParser (Maybe Step)
pFinalAnswer skipped e = do
  pLog "pFinalAnswer"
  rest <- many skip
  let answers = catMaybes $ map mAnswer (skipped ++ rest)
  let me = closestInList (filter (\n -> isNat n && n >= 12 && n <= 20) $ e : answers) 16
  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
    ]