{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}

{-|
Module      : Framework.Lambda.Signature
Description : Signatures for constants used across the framework.
Copyright   : (c) Julian Grove and Aaron Steven White, 2025
License     : MIT
Maintainer  : julian.grove@gmail.com

Signatures for constants used across the framework.
-}

module Framework.Lambda.Signature ( mkStackSig
                                  , mkStateSig
                                  , tau0
                                  , tauArithmetic
                                  , tauIndicesInit
                                  , tauLogical
                                  , tauProbProg
                                  , tauReals
                                  , tauStatesInit
                                  ) where

import Framework.Lambda.Convenience
import Framework.Lambda.Terms
import Framework.Lambda.Types

-- ** Signatures

-- | Combined signature.
tau0 :: Sig
tau0 :: Sig
tau0 = Sig
tauArithmetic  Sig -> Sig -> Sig
forall (m :: * -> *) a b.
Alternative m =>
(a -> m b) -> (a -> m b) -> a -> m b
<||>
       Sig
tauIndicesInit Sig -> Sig -> Sig
forall (m :: * -> *) a b.
Alternative m =>
(a -> m b) -> (a -> m b) -> a -> m b
<||>
       Sig
tauLogical     Sig -> Sig -> Sig
forall (m :: * -> *) a b.
Alternative m =>
(a -> m b) -> (a -> m b) -> a -> m b
<||>
       Sig
tauProbProg    Sig -> Sig -> Sig
forall (m :: * -> *) a b.
Alternative m =>
(a -> m b) -> (a -> m b) -> a -> m b
<||>
       Sig
tauReals       Sig -> Sig -> Sig
forall (m :: * -> *) a b.
Alternative m =>
(a -> m b) -> (a -> m b) -> a -> m b
<||>
       Sig
tauStatesInit

-- | Some arithmetic operators.
tauArithmetic :: Sig
tauArithmetic :: Sig
tauArithmetic = \case
  Left String
"mult" -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:× Type
r Type -> Type -> Type
:→ Type
r)
  Left String
"add"  -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:× Type
r Type -> Type -> Type
:→ Type
r)
  Left String
"max"  -> Type -> Maybe Type
forall a. a -> Maybe a
Just ((Type
r Type -> Type -> Type
:→ Type
t) Type -> Type -> Type
:→ Type
r)
  Left String
"neg"  -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:→ Type
r)
  Left String
"(≥)"  -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:→ Type
r Type -> Type -> Type
:→ Type
t)
  Either String Double
_           -> Maybe Type
forall a. Maybe a
Nothing

-- | Initial indices.
tauIndicesInit :: Sig
tauIndicesInit :: Sig
tauIndicesInit = \case
  Left String
"@" -> Type -> Maybe Type
forall a. a -> Maybe a
Just Type
σ -- Blank starting world.
  Either String Double
_        -> Maybe Type
forall a. Maybe a
Nothing

-- | Logical constants and operators defind thereon.
tauLogical :: Sig
tauLogical :: Sig
tauLogical = \case
  Left String
"∀"            -> Type -> Maybe Type
forall a. a -> Maybe a
Just ((Type
α Type -> Type -> Type
:→ Type
t) Type -> Type -> Type
:→ Type
t)
  Left String
"∃"            -> Type -> Maybe Type
forall a. a -> Maybe a
Just ((Type
α Type -> Type -> Type
:→ Type
t) Type -> Type -> Type
:→ Type
t)
  Left String
"(∧)"          -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
t Type -> Type -> Type
:→ Type
t Type -> Type -> Type
:→ Type
t)
  Left String
"(∨)"          -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
t Type -> Type -> Type
:→ Type
t Type -> Type -> Type
:→ Type
t)
  Left String
"(⇒)"          -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
t Type -> Type -> Type
:→ Type
t Type -> Type -> Type
:→ Type
t)
  Left String
"¬"            -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
t Type -> Type -> Type
:→ Type
t)
  Left String
"T"            -> Type -> Maybe Type
forall a. a -> Maybe a
Just Type
t
  Left String
"F"            -> Type -> Maybe Type
forall a. a -> Maybe a
Just Type
t
  Left String
"if_then_else" -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
t Type -> Type -> Type
:× Type
α Type -> Type -> Type
:× Type
α Type -> Type -> Type
:→ Type
α)
  Left String
"(=)"          -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
α Type -> Type -> Type
:× Type
α Type -> Type -> Type
:→ Type
t)
  Either String Double
_                   -> Maybe Type
forall a. Maybe a
Nothing

-- | Various useful probability distributions and operators defined thereon.
tauProbProg :: Sig
tauProbProg :: Sig
tauProbProg = \case
  Left String
"Bernoulli"    -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:→ Type -> Type
P Type
t)
  Left String
"Beta"         -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:× Type
r Type -> Type -> Type
:→ Type -> Type
P Type
r)
  Left String
"disj"         -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:× Type -> Type
P Type
α Type -> Type -> Type
:× Type -> Type
P Type
α Type -> Type -> Type
:→ Type -> Type
P Type
α)
  Left String
"Normal"       -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:× Type
r Type -> Type -> Type
:→ Type -> Type
P Type
r)
  Left String
"Normal_cdf"   -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:× Type
r Type -> Type -> Type
:→ Type
r Type -> Type -> Type
:→ Type
r)
  Left String
"Logit_normal" -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:× Type
r Type -> Type -> Type
:→ Type -> Type
P Type
r)
  Left String
"observe"      -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
t Type -> Type -> Type
:→ Type -> Type
P Type
Unit)
  Left String
"factor"       -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:→ Type -> Type
P Type
Unit)
  Left String
"Truncate"     -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
r Type -> Type -> Type
:× Type
r Type -> Type -> Type
:→ Type -> Type
P Type
r Type -> Type -> Type
:→ Type -> Type
P Type
r)
  Left String
"#"            -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type -> Type
P Type
α) -- The undefined distribution.
  Left String
"𝟙"            -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
t Type -> Type -> Type
:→ Type
r)
  Left String
"𝔼"            -> Type -> Maybe Type
forall a. a -> Maybe a
Just ((Type
α Type -> Type -> Type
:→ Type
r) Type -> Type -> Type
:→ Type -> Type
P Type
α Type -> Type -> Type
:→ Type
r)
  Left String
"Pr"           -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type -> Type
P Type
t Type -> Type -> Type
:→ Type
r)
  Either String Double
_                   -> Maybe Type
forall a. Maybe a
Nothing

-- | Real numbers.
tauReals :: Sig
tauReals :: Sig
tauReals = \case
  Right Double
_ -> Type -> Maybe Type
forall a. a -> Maybe a
Just Type
r
  Either String Double
_       -> Maybe Type
forall a. Maybe a
Nothing

-- | Initial states.
tauStatesInit :: Sig
tauStatesInit :: Sig
tauStatesInit = \case
  Left String
"ϵ" -> Type -> Maybe Type
forall a. a -> Maybe a
Just Type
σ -- Blank starting state.
  Either String Double
_        -> Maybe Type
forall a. Maybe a
Nothing 


-- * Functions for generating signatures

-- | Make signatures for updating and accessing states.
mkStateSig :: Type -> [(String, Type)] -> Sig
mkStateSig :: Type -> [(String, Type)] -> Sig
mkStateSig Type
_   []               = Maybe Type -> Sig
forall a b. a -> b -> a
const Maybe Type
forall a. Maybe a
Nothing
mkStateSig Type
sTy ((String
c, Type
ty) : [(String, Type)]
ctys) = Sig
ctySig Sig -> Sig -> Sig
forall (m :: * -> *) a b.
Alternative m =>
(a -> m b) -> (a -> m b) -> a -> m b
<||> Type -> [(String, Type)] -> Sig
mkStateSig Type
sTy [(String, Type)]
ctys
  where ctySig :: Sig
        ctySig :: Sig
ctySig = \case
          Left String
c' | String
c' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
c           -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
sTy Type -> Type -> Type
:→ Type
ty)
          Left String
c' | String
c' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"upd_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
c -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
ty Type -> Type -> Type
:→ Type
sTy Type -> Type -> Type
:→ Type
sTy)
          Either String Double
_                           -> Maybe Type
forall a. Maybe a
Nothing

-- | Make signatures for pushing to and popping from stacks.
mkStackSig :: Type -> (Type -> Type) -> (Type -> Type) -> [(String, Type)] -> Sig
mkStackSig :: Type -> (Type -> Type) -> (Type -> Type) -> [(String, Type)] -> Sig
mkStackSig Type
_   Type -> Type
_ Type -> Type
_ []               = Maybe Type -> Sig
forall a b. a -> b -> a
const Maybe Type
forall a. Maybe a
Nothing
mkStackSig Type
sTy Type -> Type
f Type -> Type
g ((String
c, Type
ty) : [(String, Type)]
ctys) = Sig
ctySig Sig -> Sig -> Sig
forall (m :: * -> *) a b.
Alternative m =>
(a -> m b) -> (a -> m b) -> a -> m b
<||> Type -> (Type -> Type) -> (Type -> Type) -> [(String, Type)] -> Sig
mkStackSig Type
sTy Type -> Type
f Type -> Type
g [(String, Type)]
ctys
  where ctySig :: Sig
        ctySig :: Sig
ctySig = \case
          Left String
c' | String
c' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"push_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
c -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
ty Type -> Type -> Type
:→ Type
sTy Type -> Type -> Type
:→ Type -> Type
f Type
sTy)
          Left String
c' | String
c' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"pop_"  String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
c -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type
sTy Type -> Type -> Type
:→ Type
ty Type -> Type -> Type
:× Type -> Type
g Type
sTy)
          Either String Double
_                            -> Maybe Type
forall a. Maybe a
Nothing