content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DeriveLift #-}
  {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE InstanceSigs #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE PatternSynonyms #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeFamilies #-}
  
  -- |
  -- Module      :   Grisette.Internal.SymPrim.Prim.Model
  -- Copyright   :   (c) Sirui Lu 2021-2024
  -- License     :   BSD-3-Clause (see the LICENSE file)
  --
  -- Maintainer  :   siruilu@cs.washington.edu
  -- Stability   :   Experimental
  -- Portability :   GHC only
  module Grisette.Internal.SymPrim.Prim.Model
    ( SymbolSet (..),
      ConstantSymbolSet,
      AnySymbolSet,
      Model (..),
      ModelValuePair (..),
      equation,
      evalTerm,
    )
  where
  
  import Control.DeepSeq (NFData)
  import qualified Data.Binary as Binary
  import Data.Bytes.Serial (Serial (deserialize, serialize))
  import qualified Data.HashMap.Strict as M
  import qualified Data.HashSet as S
  import Data.Hashable (Hashable)
  import Data.List (sort, sortOn)
  import Data.Proxy (Proxy (Proxy))
  import qualified Data.Serialize as Cereal
  import GHC.Generics (Generic)
  import Grisette.Internal.Core.Data.Class.ModelOps
    ( ModelOps
        ( emptyModel,
          exceptFor,
          exceptFor',
          extendTo,
          insertValue,
          isEmptyModel,
          modelContains,
          restrictTo,
          valueOf
        ),
      ModelRep (buildModel),
      SymbolSetOps
        ( containsSymbol,
          differenceSet,
          emptySet,
          insertSymbol,
          intersectionSet,
          isEmptySet,
          unionSet
        ),
      SymbolSetRep (buildSymbolSet),
    )
  import Grisette.Internal.SymPrim.GeneralFun (generalSubstSomeTerm)
  import Grisette.Internal.SymPrim.Prim.Internal.Term
    ( SomeTypedAnySymbol,
      SomeTypedConstantSymbol,
      SupportedPrim,
      SymbolKind (AnyKind, ConstantKind),
      Term,
      TypedAnySymbol,
    )
  import Grisette.Internal.SymPrim.Prim.Term
    ( IsSymbolKind,
      ModelValue,
      SomeTypedSymbol (SomeTypedSymbol),
      SupportedPrim (defaultValue),
      TypedSymbol,
      conTerm,
      defaultValueDynamic,
      pevalEqTerm,
      showUntyped,
      someTypedSymbol,
      symTerm,
      toModelValue,
      unsafeFromModelValue,
      pattern SupportedTypedSymbol,
    )
  import Language.Haskell.TH.Syntax (Lift)
  
  -- $setup
  -- >>> import Grisette.Core
  -- >>> import Grisette.SymPrim
  
  -- | Set of symbols.
  --
  -- Check 'Grisette.Core.SymbolSetOps' for operations, and
  -- 'Grisette.Core.SymbolSetRep' for manual constructions.
  newtype SymbolSet knd = SymbolSet
    { unSymbolSet :: S.HashSet (SomeTypedSymbol knd)
    }
    deriving (Eq, Generic)
    deriving newtype (Hashable)
    deriving anyclass (Serial)
  
  instance (IsSymbolKind knd) => Cereal.Serialize (SymbolSet knd) where
    put = serialize
    get = deserialize
  
  instance (IsSymbolKind knd) => Binary.Binary (SymbolSet knd) where
    put = serialize
    get = deserialize
  
  -- | Set of constant symbols. Excluding unintepreted functions.
  type ConstantSymbolSet = SymbolSet 'ConstantKind
  
  -- | Set of any symbols.
  type AnySymbolSet = SymbolSet 'AnyKind
  
  instance Semigroup (SymbolSet knd) where
    SymbolSet s1 <> SymbolSet s2 = SymbolSet $ S.union s1 s2
  
  instance Monoid (SymbolSet knd) where
    mempty = emptySet
  
  instance Show (SymbolSet knd) where
    showsPrec prec (SymbolSet s) = showParen (prec >= 10) $ \x ->
      "SymbolSet {"
        ++ go0 (sort $ show <$> S.toList s)
        ++ "}"
        ++ x
      where
        go0 [] = ""
        go0 [x] = x
        go0 (x : xs) = x ++ ", " ++ go0 xs
  
  -- | Model returned by the solver.
  --
  -- Check 'Grisette.Core.ModelOps' for operations, and 'Grisette.Core.ModelRep'
  -- for manual constructions.
  newtype Model = Model
    { unModel :: M.HashMap SomeTypedAnySymbol ModelValue
    }
    deriving stock (Eq, Generic, Lift)
    deriving newtype (Hashable, NFData)
    deriving anyclass (Serial)
  
  instance Cereal.Serialize Model where
    put = serialize
    get = deserialize
  
  instance Binary.Binary Model where
    put = serialize
    get = deserialize
  
  instance Semigroup Model where
    Model m1 <> Model m2 = Model $ M.union m1 m2
  
  instance Monoid Model where
    mempty = emptyModel
  
  instance Show Model where
    showsPrec prec (Model m) = showParen (prec >= 10) $ \x ->
      "Model {"
        ++ go0 (sortOn (\(x, _) -> show x) $ M.toList m)
        ++ "}"
        ++ x
      where
        go0 [] = ""
        go0 [(SomeTypedSymbol s, v)] = showUntyped s ++ " -> " ++ show v
        go0 ((SomeTypedSymbol s, v) : xs) = showUntyped s ++ " -> " ++ show v ++ ", " ++ go0 xs
  
  -- | Given a typed symbol and a model, return the equation (symbol = value)
  -- encoded in the model.
  equation :: TypedAnySymbol a -> Model -> Maybe (Term Bool)
  equation tsym@SupportedTypedSymbol m =
    case valueOf tsym m of
      Just v -> Just $ pevalEqTerm (symTerm tsym) (conTerm v)
      Nothing -> Nothing
  
  instance SymbolSetOps (SymbolSet knd) (TypedSymbol knd) where
    emptySet = SymbolSet S.empty
    isEmptySet (SymbolSet s) = S.null s
    containsSymbol s =
      S.member (someTypedSymbol s) . unSymbolSet
    insertSymbol s = SymbolSet . S.insert (someTypedSymbol s) . unSymbolSet
    intersectionSet (SymbolSet s1) (SymbolSet s2) = SymbolSet $ S.intersection s1 s2
    unionSet (SymbolSet s1) (SymbolSet s2) = SymbolSet $ S.union s1 s2
    differenceSet (SymbolSet s1) (SymbolSet s2) = SymbolSet $ S.difference s1 s2
  
  instance
    SymbolSetRep
      (SomeTypedSymbol knd)
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet sym = SymbolSet $ S.singleton sym
  
  instance
    SymbolSetRep
      [SomeTypedSymbol knd]
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet = SymbolSet . S.fromList
  
  instance
    SymbolSetRep
      [TypedSymbol knd t]
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet sym = buildSymbolSet $ someTypedSymbol <$> sym
  
  instance
    SymbolSetRep
      (TypedSymbol knd t)
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet sym = insertSymbol sym emptySet
  
  instance
    SymbolSetRep
      ( TypedSymbol knd a,
        TypedSymbol knd b
      )
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet (sym1, sym2) =
      insertSymbol sym2
        . insertSymbol sym1
        $ emptySet
  
  instance
    SymbolSetRep
      ( TypedSymbol knd a,
        TypedSymbol knd b,
        TypedSymbol knd c
      )
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet (sym1, sym2, sym3) =
      insertSymbol sym3
        . insertSymbol sym2
        . insertSymbol sym1
        $ emptySet
  
  instance
    SymbolSetRep
      ( TypedSymbol knd a,
        TypedSymbol knd b,
        TypedSymbol knd c,
        TypedSymbol knd d
      )
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet (sym1, sym2, sym3, sym4) =
      insertSymbol sym4
        . insertSymbol sym3
        . insertSymbol sym2
        . insertSymbol sym1
        $ emptySet
  
  instance
    SymbolSetRep
      ( TypedSymbol knd a,
        TypedSymbol knd b,
        TypedSymbol knd c,
        TypedSymbol knd d,
        TypedSymbol knd e
      )
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet (sym1, sym2, sym3, sym4, sym5) =
      insertSymbol sym5
        . insertSymbol sym4
        . insertSymbol sym3
        . insertSymbol sym2
        . insertSymbol sym1
        $ emptySet
  
  instance
    SymbolSetRep
      ( TypedSymbol knd a,
        TypedSymbol knd b,
        TypedSymbol knd c,
        TypedSymbol knd d,
        TypedSymbol knd e,
        TypedSymbol knd f
      )
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet (sym1, sym2, sym3, sym4, sym5, sym6) =
      insertSymbol sym6
        . insertSymbol sym5
        . insertSymbol sym4
        . insertSymbol sym3
        . insertSymbol sym2
        . insertSymbol sym1
        $ emptySet
  
  instance
    SymbolSetRep
      ( TypedSymbol knd a,
        TypedSymbol knd b,
        TypedSymbol knd c,
        TypedSymbol knd d,
        TypedSymbol knd e,
        TypedSymbol knd f,
        TypedSymbol knd g
      )
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet (sym1, sym2, sym3, sym4, sym5, sym6, sym7) =
      insertSymbol sym7
        . insertSymbol sym6
        . insertSymbol sym5
        . insertSymbol sym4
        . insertSymbol sym3
        . insertSymbol sym2
        . insertSymbol sym1
        $ emptySet
  
  instance
    SymbolSetRep
      ( TypedSymbol knd a,
        TypedSymbol knd b,
        TypedSymbol knd c,
        TypedSymbol knd d,
        TypedSymbol knd e,
        TypedSymbol knd f,
        TypedSymbol knd g,
        TypedSymbol knd h
      )
      (SymbolSet knd)
      (TypedSymbol knd)
    where
    buildSymbolSet (sym1, sym2, sym3, sym4, sym5, sym6, sym7, sym8) =
      insertSymbol sym8
        . insertSymbol sym7
        . insertSymbol sym6
        . insertSymbol sym5
        . insertSymbol sym4
        . insertSymbol sym3
        . insertSymbol sym2
        . insertSymbol sym1
        $ emptySet
  
  instance ModelOps Model AnySymbolSet TypedAnySymbol where
    emptyModel = Model M.empty
    isEmptyModel (Model m) = M.null m
    valueOf :: forall t. TypedAnySymbol t -> Model -> Maybe t
    valueOf sym@SupportedTypedSymbol (Model m) =
      (unsafeFromModelValue @t) <$> M.lookup (someTypedSymbol sym) m
    modelContains sym (Model m) = M.member (someTypedSymbol sym) m
    exceptFor (SymbolSet s) (Model m) = Model $ S.foldl' (flip M.delete) m s
    exceptFor' s (Model m) = Model $ M.delete (someTypedSymbol s) m
    restrictTo (SymbolSet s) (Model m) =
      Model $
        S.foldl'
          ( \acc sym -> case M.lookup sym m of
              Just v -> M.insert sym v acc
              Nothing -> acc
          )
          M.empty
          s
    extendTo (SymbolSet s) (Model m) =
      Model $
        S.foldl'
          ( \acc
             sym@(SomeTypedSymbol (SupportedTypedSymbol :: TypedAnySymbol t)) ->
                case M.lookup sym acc of
                  Just _ -> acc
                  Nothing -> M.insert sym (defaultValueDynamic (Proxy @t)) acc
          )
          m
          s
    insertValue sym@SupportedTypedSymbol (v :: t) (Model m) =
      Model $ M.insert (someTypedSymbol sym) (toModelValue v) m
  
  -- | Evaluate a term in the given model.
  evalTerm ::
    (SupportedPrim a) =>
    Bool ->
    Model ->
    S.HashSet SomeTypedConstantSymbol ->
    Term a ->
    Term a
  evalTerm fillDefault (Model ma) =
    generalSubstSomeTerm
      ( \(sym@SupportedTypedSymbol :: TypedSymbol 'AnyKind a) ->
          case (M.lookup (someTypedSymbol sym) ma) of
            Nothing ->
              if fillDefault
                then conTerm (defaultValue @a)
                else symTerm sym
            Just dy ->
              conTerm (unsafeFromModelValue @a dy)
      )
  
  -- |
  -- A type used for building a model by hand.
  --
  -- >>> buildModel ("x" ::= (1 :: Integer), "y" ::= True) :: Model
  -- Model {x -> 1 :: Integer, y -> true :: Bool}
  data ModelValuePair t = (TypedAnySymbol t) ::= t deriving (Show)
  
  instance ModelRep (ModelValuePair t) Model where
    buildModel (sym ::= val) = insertValue sym val emptyModel
  
  instance (ModelRep a Model, ModelRep b Model) => ModelRep (a, b) Model where
    buildModel (a, b) = buildModel a <> buildModel b
  
  instance
    ( ModelRep a Model,
      ModelRep b Model,
      ModelRep c Model
    ) =>
    ModelRep (a, b, c) Model
    where
    buildModel (a, b, c) = buildModel a <> buildModel b <> buildModel c
  
  instance
    ( ModelRep a Model,
      ModelRep b Model,
      ModelRep c Model,
      ModelRep d Model
    ) =>
    ModelRep (a, b, c, d) Model
    where
    buildModel (a, b, c, d) =
      buildModel a <> buildModel b <> buildModel c <> buildModel d
  
  instance
    ( ModelRep a Model,
      ModelRep b Model,
      ModelRep c Model,
      ModelRep d Model,
      ModelRep e Model
    ) =>
    ModelRep (a, b, c, d, e) Model
    where
    buildModel (a, b, c, d, e) =
      buildModel a <> buildModel b <> buildModel c <> buildModel d <> buildModel e
  
  instance
    ( ModelRep a Model,
      ModelRep b Model,
      ModelRep c Model,
      ModelRep d Model,
      ModelRep e Model,
      ModelRep f Model
    ) =>
    ModelRep (a, b, c, d, e, f) Model
    where
    buildModel (a, b, c, d, e, f) =
      buildModel a
        <> buildModel b
        <> buildModel c
        <> buildModel d
        <> buildModel e
        <> buildModel f
  
  instance
    ( ModelRep a Model,
      ModelRep b Model,
      ModelRep c Model,
      ModelRep d Model,
      ModelRep e Model,
      ModelRep f Model,
      ModelRep g Model
    ) =>
    ModelRep (a, b, c, d, e, f, g) Model
    where
    buildModel (a, b, c, d, e, f, g) =
      buildModel a
        <> buildModel b
        <> buildModel c
        <> buildModel d
        <> buildModel e
        <> buildModel f
        <> buildModel g
  
  instance
    ( ModelRep a Model,
      ModelRep b Model,
      ModelRep c Model,
      ModelRep d Model,
      ModelRep e Model,
      ModelRep f Model,
      ModelRep g Model,
      ModelRep h Model
    ) =>
    ModelRep (a, b, c, d, e, f, g, h) Model
    where
    buildModel (a, b, c, d, e, f, g, h) =
      buildModel a
        <> buildModel b
        <> buildModel c
        <> buildModel d
        <> buildModel e
        <> buildModel f
        <> buildModel g
        <> buildModel h
  