-----------------------------------------------------------------------------
-- 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 commonly used functions within the subexpression recognizer.
--
-----------------------------------------------------------------------------

module Recognize.SubExpr.Functions
   ( -- * Functions on Symbols
     isBuggySymbol, isBuggy, isOrSymbol, isOr, isAndSymbol, isAnd, isVarSymbol
   , isVar, hasVar, isLtSymbol, isLt, isMatchSymbol, isMatch, isMatched
   , hasMatch, extractMatched, getMatched, subMatched, isMagicVarSymbol
   , isMagicVar, hasMagicVar, isMagicNatSymbol, isMagicNat, hasMagicNat
   , isMagicNumberSymbol, isMagicNumber, hasMagicNumber
   , isLabelSymbol, isLabel, isStopSymbol, isStop
   , isSimSymbol, isSim, isSimplified, isNoSimSymbol, isNoSim, isSubSymbol
   , isSub, substituteAllIf, hasSub
     -- * Auxillary functions
   , cleanExpr, underSubst, subExprsCombs, alternativesExpr, isSubExprOf
   ) where

import Control.Arrow
import Control.Monad
import Data.List
import qualified Data.Map as M
import Data.Maybe
import Data.Tuple
import Data.Monoid ((<>))
import Domain.Math.Expr.Data
import Ideas.Common.Rewriting hiding (hasVar)
import Recognize.Data.Attribute
import Recognize.Expr.Normalform
import Recognize.SubExpr.Symbols

isBuggySymbol :: Symbol -> Bool
isBuggySymbol  = (== buggySymbol)

isBuggy :: Expr -> Bool
isBuggy (Sym s [_,_]) = isBuggySymbol s
isBuggy _ = False

isOrSymbol :: Symbol -> Bool
isOrSymbol = (== orSymbol)

isOr :: Expr -> Bool
isOr (Sym s [_,_]) = isOrSymbol s

isAndSymbol :: Symbol -> Bool
isAndSymbol = (== andSymbol)

isAnd :: Expr -> Bool
isAnd (Sym s [_,_]) = isAndSymbol s
isAnd _ = False

isVarSymbol :: Symbol -> Bool
isVarSymbol = (== varSymbol)

isVar :: Expr -> Bool
isVar (Sym s [Var _]) = isVarSymbol s
isVar _ = False

hasVar :: Expr -> Bool
hasVar e = case getFunction e of
  Nothing -> False
  Just (s,xs) -> isVarSymbol s || any hasVar xs

substituteIf :: (Expr -> Bool) -> M.Map String Expr -> Expr -> Maybe Expr
substituteIf p dic e@(Sym s [Var v])
  | isVarSymbol s && p e = M.lookup v dic
  | otherwise = Just e
substituteIf p dic e@(Var v)
  | p e = M.lookup v dic
  | otherwise = Just e
substituteIf _ _ e = Just e

substituteAllIf :: (Expr -> Bool) -> M.Map String Expr -> Expr -> Maybe Expr
substituteAllIf p dic e = case getFunction e of
  Nothing -> substituteIf p dic e
  Just (s,[Var _])
    | isVarSymbol s -> substituteIf p dic e
    | otherwise     -> Just e
  Just (s,xs) -> do
    xs' <- mapM (substituteAllIf p dic) xs
    return $ function s xs'

-- | Builds a dictionary mapping variable strings to values
-- these values are calculated by applying a function to one of those strings
mkVarMap :: Expr -> M.Map String String
mkVarMap e = case getFunction e of
  Nothing -> M.empty
  Just (_,[Var s]) -> M.singleton s s
  Just (_,xs) -> foldr (\a dic -> mkVarMap a <> dic) M.empty xs

isLtSymbol :: Symbol -> Bool
isLtSymbol = (== ltSymbol)

isLt :: Expr -> Bool
isLt (Sym s _) = isLtSymbol s
isLt _ = False

isMatchSymbol :: Symbol -> Bool
isMatchSymbol = (== matchSymbol)

isMatch :: Expr -> Bool
isMatch (Sym s [_]) = isMatchSymbol s
isMatch _ = False

isMatched :: Expr -> Bool
isMatched e = case getFunction e of
  Just (s,[x]) -> isMatchSymbol s || (isNoSimSymbol s && isMatched x)
  Just (s,[x,y]) -> (isAndSymbol s && isMatched x && isMatched y)
    || (isOrSymbol s && (isMatched x || isMatched y))
    || (isLabelSymbol s && isMatched y)
  Just (s,[_,_,z])
    | isLabelSymbol s -> isMatched z
  _ -> False

-- | Does this expression have a matched subexpression?
hasMatch :: Expr -> Bool
hasMatch e = case getFunction e of
  Nothing -> False
  Just (s,xs)
    | isMatchSymbol s -> True
    | otherwise -> any hasMatch xs

-- | If this expression has been matched, then extracted the matched expression
extractMatched :: Expr -> Expr
extractMatched e@(Sym s [x])
  | isMatchSymbol s = x
  | otherwise = e
extractMatched e = e

getMatched :: Expr -> Maybe Expr
getMatched e = case getFunction e of
  Nothing -> Nothing
  Just (s,[x])
    | isMatchSymbol s -> Just x
    | otherwise -> getMatched x
  Just (s,[x,y])
    | isLabelSymbol s -> getMatched y
    | isOrSymbol s -> msum $ map getMatched [x,y]
    | otherwise -> do
      x' <- getMatched x
      y' <- getMatched y
      return $ binary s x' y'
  Just (s,[_,_,z])
    | isLabelSymbol s -> getMatched z
  Just (s,xs) -> do
      xs' <- mapM getMatched xs
      return $ function s xs'

-- | Substitue any matching with its matched expression
subMatched :: Expr -> Expr
subMatched e = case getFunction e of
  Nothing -> e
  Just (s,[x])
    | isMatchSymbol s -> x
    | otherwise -> unary s $ subMatched x
  Just (s,xs) -> function s $ map subMatched xs

-- Matches to any variable
isMagicVarSymbol :: Symbol -> Bool
isMagicVarSymbol = (== magicVarSymbol)

isMagicVar :: Expr -> Bool
isMagicVar (Sym s _) = isMagicVarSymbol s
isMagicVar _ = False

hasMagicVar :: Expr -> Bool
hasMagicVar e = case getFunction e of
  Nothing -> False
  Just (s,xs) -> isMagicVarSymbol s || any isMagicVar xs || any hasMagicVar xs

-- Matches to any natural number
isMagicNatSymbol :: Symbol -> Bool
isMagicNatSymbol = (== magicNatSymbol)

isMagicNat :: Expr -> Bool
isMagicNat (Sym s []) = isMagicNatSymbol s
isMagicNat _ = False

hasMagicNat :: Expr -> Bool
hasMagicNat e = case getFunction e of
  Nothing -> False
  Just (s,xs) -> isMagicNatSymbol s || any isMagicNat xs || any hasMagicNat xs

-- Matches to any natural number
isMagicNumberSymbol :: Symbol -> Bool
isMagicNumberSymbol = (== magicNumberSymbol)

isMagicNumber :: Expr -> Bool
isMagicNumber (Sym s []) = isMagicNumberSymbol s
isMagicNumber _ = False

hasMagicNumber :: Expr -> Bool
hasMagicNumber e = case getFunction e of
  Nothing -> False
  Just (s,xs) -> isMagicNumberSymbol s || any isMagicNumber xs || any hasMagicNumber xs

isLabelSymbol :: Symbol -> Bool
isLabelSymbol = (== labelSymbol)

isLabel :: Expr -> Bool
isLabel (Sym s [_,_]) = isLabelSymbol s
isLabel (Sym s [_,_,_]) = isLabelSymbol s
isLabel _ = False

isStopSymbol :: Symbol -> Bool
isStopSymbol = (== stopSymbol)

isStop :: Expr -> Bool
isStop (Sym s [_]) = isStopSymbol s
isStop _ = False

isSimSymbol :: Symbol -> Bool
isSimSymbol = (== simSymbol)

isSim :: Expr -> Bool
isSim (Sym s [_]) = isSimSymbol s
isSim _ = False

isSimplified :: Expr -> Bool
isSimplified e = cleanExpr e == subMatched e && nf e == e

isNoSimSymbol :: Symbol -> Bool
isNoSimSymbol = (== noSimSymbol)

isNoSim :: Expr -> Bool
isNoSim (Sym s [_]) = isNoSimSymbol s
isNoSim _ = False

isSubSymbol :: Symbol -> Bool
isSubSymbol = (== subSymbol)

isSub :: Expr -> Bool
isSub (Sym s [_]) = isSubSymbol s
isSub _ = False

hasSub :: Expr -> Bool
hasSub e = case getFunction e of
  Nothing -> False
  Just (s,xs) -> isSubSymbol s || any hasSub xs

-- | Remove all custom symbols from the expression
-- Useful/necessary for comandisons.
-- Note that it is not possible to remove all symbols (@and@,or,magicNat,magicVar)
cleanExpr :: Expr -> Expr
cleanExpr e = case getFunction e of
  Nothing -> e
  Just (s,[x])
    | isMatchSymbol s -> cleanExpr x
    | isSimSymbol s -> cleanExpr x
    | isNoSimSymbol s -> cleanExpr x
    | isSubSymbol s -> cleanExpr x
    | isStopSymbol s -> cleanExpr x
    | otherwise -> unary s (cleanExpr x)
  Just (s,[x,y])
    | isLabelSymbol s -> cleanExpr y
    | isAndSymbol s -> cleanExpr x <&> cleanExpr y
    | isOrSymbol s -> cleanExpr x <?> cleanExpr y
    | otherwise -> binary s (cleanExpr x) (cleanExpr y)
  Just (s,[_,_,z])
    | isLabelSymbol s -> cleanExpr z
  Just (s,xs) -> function s (map cleanExpr xs)

-- | Substitutes all special vars with a unique regular Var
-- then applies some function after which the substitution is reversed
underSubst :: (Expr -> (Expr,[Attribute])) -> Expr -> Maybe (Expr, [Attribute])
underSubst f e = do
  let dic = mkVarMap e
  -- M.Map String Expr, actually just M.Map String (Var v)
  let dicU = snd $ M.mapAccum (\acc _ -> (acc+1, Var $ "$$$" ++ show acc)) 0 dic
  eU <- substituteAllIf isVar dicU e
  -- M.Map String String, this is partial but we know that dicU only contains Vars
  let dicU' = M.map (\(Var v) -> v) dicU
  -- M.Map String String
  dicUInverse <- invert dicU'
  -- M.Map String Expr, actually just M.Map String (Sym varSymbol [Var v])
  let dicUInverse' = M.map var dicUInverse
  -- Now we have a simplified form, but we must still replace every var with the expressions
  -- that they refer to.
  let (fe, attr) = f eU
  -- Only substitute a variable if it is present in the dicUInverse' Map, since there might be
  -- other variables already present in the expression
  fe' <- substituteAllIf (\e -> isVariable e &&
    case e of
      Var v -> isJust $ M.lookup v dicUInverse') dicUInverse' fe
  -- TODO: Added attributes now contain the substitute value introduced here instead of the corresponding subexpression variable
  return (fe',attr)

-- | Compute all possible combinations for a list of expressions
subExprsCombs :: [Expr] -> [([Expr],[Attribute])]
subExprsCombs [] = [([],[])]
subExprsCombs (x:xs) =
  let xs' = subExprsCombs xs
  in concatMap (\(e, rw) -> map ((:) e *** (++) rw) xs') (alternativesExpr x)

alternativesExpr :: Expr -> [(Expr, [Attribute])]
alternativesExpr e = case getFunction e of
  Nothing -> [(e,[])]
  Just (s,[x])
    | isSimSymbol s -> [(nf x,[])]
    | otherwise -> first (unary s) <$> alternativesExpr x

  Just (s,[x,y])
    | isBuggySymbol s -> (second (CommonMistake:) <$> alternativesExpr y) ++ alternativesExpr x
    | isOrSymbol s || isAndSymbol s -> alternativesExpr x ++ alternativesExpr y
  Just (s,xs) -> map (first (function s)) $ nary xs
  _ -> error (show e)
  where
    nary :: [Expr] -> [([Expr],[Attribute])]
    nary [] = [([],[])]
    nary (x:xs) = do
      (xe,xa) <- alternativesExpr x
      (xes,xas) <- nary xs
      return (xe:xes,xa++xas)



isSubExprOf :: Expr -> Expr -> Bool
isSubExprOf x e = x == e || (case getFunction e of
  Nothing -> False
  Just (s,xs) -> any (isSubExprOf x) xs)

invert :: Ord v => M.Map k v -> Maybe (M.Map v k)
invert m = do
  let kvs = M.toList m
  let vs = map snd kvs
  guard (length (nub vs) == length vs)
  return $ M.fromList $ map swap kvs