-----------------------------------------------------------------------------
-- 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.ComponentSet
  ( -- * Types
    ComponentSet, module Domain.Statistics.Component
    -- * Constructors
  , initialSet, derivedSet
    -- * Combining and transforming
  , append, delete, mapComponent, initials, derived, toList
   -- * Getters
  , get, getData, getExpr, getRelation, getRhsExpr
  , getTestType, getSided, getRejectionHypotheses, getConclusion
    -- * Membership
  , contains, doesNotContain
  , isInitial, isDerived
  ) where

import Control.Monad
import Data.List hiding (delete)
import Data.Maybe
import Domain.Statistics.Symbols
import Domain.Math.Data.Relation
import Domain.Math.Expr
import Domain.Statistics.Data
import Domain.Statistics.Component
import Ideas.Common.Rewriting hiding (trueSymbol,falseSymbol)
import Ideas.Common.Classes

------------------------------------------------------------------------------
-- Types

data ComponentSet = CS 
   { initialsList :: [(ComponentId, Component)]
   , derivedList  :: [(ComponentId, Component)]
   }

instance Show ComponentSet where
   show cs = unlines (map (f True) (initialsList cs) ++ map (f False) (derivedList cs))
    where
      f b (k, v) =  bracketsIf b (getCId k) ++ ": " ++ show v ++ ";"

      bracketsIf True  s = "[" ++ s ++ "]"
      bracketsIf False s = s

instance Eq ComponentSet where
   CS xs1 xs2 == CS ys1 ys2 = f xs1 == f ys1 && f xs2 == f ys2
    where
      f = sortBy (\(x, _) (y, _) -> x `compare` y)

instance Semigroup ComponentSet where
   CS xs1 xs2 <> CS ys1 ys2 = CS (xs1 ++ ys1) (xs2 ++ ys2)

instance Monoid ComponentSet where
   mempty  = CS [] []
   mappend = (<>)

instance IsTerm ComponentSet where
   toTerm cs = TList $ 
      [ ternary componentSymbol (toTerm k) (symbol initialSymbol) (toTerm c) 
      | (k, c) <- initialsList cs
      ] ++
      [ ternary componentSymbol (toTerm k) (symbol derivedSymbol) (toTerm c) 
      | (k, c) <- derivedList cs
      ]

   fromTerm (TList xs) =
      let f (TCon s [t1, TCon sv [], t3]) | s == componentSymbol = do
             k  <- fromTerm t1
             c  <- fromTerm t3
             if sv == initialSymbol 
                then return $ CS [(k, c)] []
                else if sv == derivedSymbol 
                then return $ CS [] [(k, c)]
                else fail "expected initial/derived"
          f _ = fail "Term is not a component"
      in mconcat <$> mapM f xs
   fromTerm _ = fail "Term is not a component (not a list)"

------------------------------------------------------------------------------
-- Constructors

initialSet :: [(ComponentId, Component)] -> ComponentSet
initialSet xs = CS xs []

derivedSet :: [(ComponentId, Component)] -> ComponentSet
derivedSet = CS []

------------------------------------------------------------------------------
-- Combining and transforming

append :: ComponentId -> Component -> ComponentSet -> ComponentSet
append k v cs = cs <> derivedSet [(k, v)]

delete :: ComponentId -> ComponentSet -> ComponentSet
delete k (CS xs ys) = CS (f xs) (f ys)
 where
   f = filter ((/= k) . fst)

mapComponent :: (Component -> Component) -> ComponentSet -> ComponentSet
mapComponent f (CS xs ys) = CS (change xs) (change ys)
 where
   change = map (mapSecond f)
   
initials :: ComponentSet -> ComponentSet
initials cs = mempty { initialsList = initialsList cs }

derived :: ComponentSet -> ComponentSet 
derived cs = mempty { derivedList = derivedList cs }

toList :: ComponentSet -> [(ComponentId, Component)]
toList cs = initialsList cs ++ derivedList cs

------------------------------------------------------------------------------
-- Getters

get :: Monad m => ComponentId -> ComponentSet -> m Component
get n cs = 
   -- first consider derived components
   case lookup n (derivedList cs ++ initialsList cs) of
      Just c  -> return c
      Nothing -> fail $ show n ++ " missing"

getData :: Monad m => ComponentId -> ComponentSet -> m Data
getData n = get n >=> isData

getExpr :: Monad m => ComponentId -> ComponentSet -> m Expr
getExpr n = get n >=> isExpr

getRelation :: Monad m => ComponentId -> ComponentSet -> m (Relation Expr)
getRelation n = get n >=> isRelation

getRhsExpr :: MonadPlus m => ComponentId -> ComponentSet -> m Expr
getRhsExpr n cs = fmap rightHandSide (getRelation n cs) `mplus` getExpr n cs

getTestType :: Monad m => ComponentId -> ComponentSet -> m TestType
getTestType n = get n >=> isTestType

getSided :: Monad m => ComponentId -> ComponentSet -> m Sided
getSided n = get n >=> isSided

getConclusion :: Monad m => ComponentId -> ComponentSet -> m Bool
getConclusion n = get n >=> isConclusion

getRejectionHypotheses :: Monad m => ComponentId -> ComponentSet -> m RejectionHypotheses
getRejectionHypotheses n = get n >=> isRejectionHypotheses

------------------------------------------------------------------------------
-- Membership

contains :: ComponentSet -> ComponentId -> Bool
contains cs x = isJust (get x cs)

doesNotContain :: ComponentSet -> ComponentId -> Bool
doesNotContain cs = not . contains cs

isInitial :: ComponentId -> ComponentSet -> Bool
isInitial n cs = initials cs `contains` n

isDerived :: ComponentId -> ComponentSet -> Bool
isDerived n cs = derived cs `contains` n