-----------------------------------------------------------------------------
-- Copyright 2020, Ideas 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.
-----------------------------------------------------------------------------
module Domain.Statistics.Views
   ( Substitution, Substitutable(..), single
   , getSubstitution, substitute
   , checkSubstitution
   , substitutedView, evaluatedView
   ) where

import qualified Data.Map as M
import qualified Data.Set as S
import Data.List (intercalate)
import Data.Maybe
import Domain.Math.Data.Relation
import Domain.Math.Expr
import Domain.Math.Numeric.Views
import Domain.Statistics.ComponentSet
import Ideas.Common.View
import Ideas.Utils.Uniplate

-- idempotent, non-recursive substitutions
newtype Substitution = S (M.Map String Expr)

instance Show Substitution where
   show (S m) = "{" ++ intercalate ", " (map f (M.toList m)) ++ "}"
    where
      f (s, expr) = s ++ "=" ++ show expr

instance Semigroup Substitution where
   (<>) (S m) = flip (M.foldrWithKey extend) m

instance Monoid Substitution where
   mempty  = S M.empty
   mappend = (<>)

member :: String -> Substitution -> Bool
member s (S m) = s `M.member` m

single :: String -> Expr -> Substitution
single s expr 
   | s `elem` vars expr = mempty
   | otherwise          = S (M.singleton s expr)

extend :: String -> Expr -> Substitution -> Substitution
extend s expr sub@(S m)
   | s `elem` vars expr = mempty -- recursive expression
   | otherwise =
        case M.lookup s m of
           Nothing -> S (M.insert s (sub |-> expr) (M.map (single s expr |->) m))
           Just e 
              | e == (sub |-> expr) -> sub 
                -- inconsistent extension; should not happen.
                -- deleting s in m breaks monoid property for Substitution
              | otherwise -> S (M.delete s m) 

-- throws an error for non-idempotent substitutions
checkSubstitution :: Substitution -> Substitution 
checkSubstitution sub@(S m)
   | disjoint (M.keysSet m) (vars (M.elems m)) = sub 
   | otherwise = error $ "Non-idempotent substitution: " ++ show sub
 where
   -- in Data.Set, since 0.5.11
   disjoint x y = S.null (S.intersection x y)

class Substitutable a where 
   (|->) :: Substitution -> a -> a
   vars  :: a -> S.Set String

instance Substitutable a => Substitutable [a] where
   (|->) = map . (|->)
   vars  = S.unions . map vars

instance Substitutable Expr where
   S xs |-> Var s = fromMaybe (Var s) (M.lookup s xs)
   sub  |-> expr  = descend (sub |->) expr
   
   vars (Var s)   = S.singleton s
   vars expr      = S.unions (map vars (children expr))

instance Substitutable a => Substitutable (Relation a) where
   sub |-> r = fmap (sub |->) r
   vars r    = vars (leftHandSide r) `S.union` vars (rightHandSide r)

instance Substitutable Component where
   sub |-> CExpr e     = CExpr $ sub |-> e
   sub |-> CRelation r = CRelation $ sub |-> r
   _   |-> c           = c
   
   vars (CExpr e)     = vars e
   vars (CRelation r) = vars r
   vars _             = S.empty

instance Substitutable ComponentSet where
   sub |-> cs = mapComponent (sub |->) cs
   vars = S.unions . map (vars . snd) . toList

getSubstitution :: ComponentSet -> Substitution
getSubstitution = addSigmaM . addSEM . make . resolveTestFormulaValue . addComponentVars
 where 
   -- if test value and test formula are both present (but differ because of 
   -- rounding), prefer the test value
   resolveTestFormulaValue cs 
      | cs `contains` TestValue && cs `contains` TestFormula = 
           delete TestFormula cs
      | otherwise = cs
 
   make cs = mconcat (map f (xs ++ ys))
    where
      -- prefer initials in substitution
      xs = toList (initials cs)
      ys = filter notInital (toList (derived cs))
      notInital = (`notElem` map fst xs) . fst
    
      f (cid, CExpr e) = single (getCId cid) e
      f (_, CRelation r) | relationType r == EqualTo =
         case leftHandSide r of
            Var v -> single v (rightHandSide r)
            _ -> mempty
      f _ = mempty
   
   -- add formulas, but only if not already present in the substitution
   addSigmaM sub
      |  "sigmaM" `member` sub = sub
      | otherwise = sub <> single "sigmaM" (Var "sigma" / sqrt (Var "n"))
   addSEM sub 
      | "SEM" `member` sub = sub
      | otherwise = sub <> single "SEM" (Var "s" / sqrt (Var "n"))

substitute :: ComponentSet -> ComponentSet
substitute cs = getSubstitution cs |-> cs

addComponentVars :: ComponentSet -> ComponentSet
addComponentVars cs = foldr (uncurry addComponentVar) (addPopulationMean cs) table
 where
   table = 
      [ (PopulationSdev, "sigma")
      , (SampleSdev, "s")
      , (SampleMean, "M")
      , (One SampleMean, "M1")
      , (Two SampleMean, "M2")
      , (One SampleSdev, "s1")
      , (Two SampleSdev, "s2")
      , (PopulationMean, "mu")
      ]

-- extract population mean from alternative hypothesis
addPopulationMean :: ComponentSet -> ComponentSet
addPopulationMean cs =
   case (get PopulationMean cs, getRelation AlternativeHypothesis cs) of
      (Nothing, Just rel) -> append PopulationMean (CExpr (rightHandSide rel)) cs
      _ -> cs

addComponentVar :: ComponentId -> String -> ComponentSet -> ComponentSet
addComponentVar cid var cs =
   case getRhsExpr cid cs of
      Just e | cs `doesNotContain` Other var ->
         -- to do: is it necessary to add (Other var) to the component set??
         append (Other var) (CExpr e) (single var e |-> cs)
      _ -> cs

----------------------------------------------------------
-- Substituted view

substitutedView :: View ComponentSet ComponentSet
substitutedView = makeView (Just . substitute) id

----------------------------------------------------------
-- Evaluated view

evaluatedView :: View ComponentSet ComponentSet
evaluatedView = makeView f id
 where
   f :: ComponentSet -> Maybe ComponentSet
   f cs = Just $ mapComponent evaluateComponent cs

evaluateComponent :: Component -> Component
evaluateComponent (CExpr x)     = CExpr $ simplify doubleView x
evaluateComponent (CRelation r) = CRelation $ fmap (simplify doubleView) r
evaluateComponent c             = c