{-# LANGUAGE GADTs #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}
{-|
Module      : Logics

The module define the logical languages and everything that is needed inside them
* agents and variables 
* expressions (boolean and integer expressions),
* programs inside the box
-}

module Logics where
import Data.Int
import Data.Map (Map, fromList, (!))
import Debug.Trace

----------------------------
-- * Agents and Variables
--------------------------
-- | an 'Agent' is determined by its identity and the set of nonobservable variables
data Agent = Agent {identity :: Int, nonobs :: [Var]}
   deriving (Eq, Ord)
data Var = BVar String
   deriving (Eq, Ord)


-- | Function 'varname' returns the name of a variable (string) 
varname :: Var -> String
varname (BVar a) = a

-- | 'intDomain' gives the type for the domain for each integer variable.
-- For now 'intDomain' is defined here, we cannot give it at each example file
-- because it is used to constrain quantifications in the "Tau" module

------------------------------------------------
-- * Expressions (boolean expressions)
------------------------------------------------
data Expr a where
    B      :: Bool -> Expr Bool        -- ^ Boolean constant 
    BValue :: Var -> Expr Bool         -- ^ Boolean variable
    BEq    :: Expr Bool -> Expr Bool -> Expr Bool
    Xor    :: [Expr Bool] -> Expr Bool  
    Or     :: [Expr Bool] -> Expr Bool
    And    :: [Expr Bool] -> Expr Bool


type BExpr = Expr Bool

-----------------------------------
-- * Logical Formulas \(\mathcal{L_K}, \mathcal{L_{\square K}}, \mathcal{L}_{FO} \) 
-----------------------------------
data Formula t where
   Atom      :: BExpr -> Formula t
   Neg       :: Formula t -> Formula t
   Conj      :: [Formula t] -> Formula t
   Disj      :: [Formula t] -> Formula t
   Imp       :: Formula t -> Formula t -> Formula t
   Equiv     :: Formula t -> Formula t -> Formula t
   K         :: Agent -> Formula Modal -> Formula Modal    -- ^ Epistemic modal operator
   -- KV        :: Agent -> Var -> Formula Modal         
   Box       :: Prog -> Formula Modal -> Formula Modal         -- ^ Dynamic modal operator
   ForAllB   :: Int ->  Formula Modal -> Formula Modal -- we could have changed Int into Str or even Var, but that is not necessary
   ExistsB   :: Int ->  Formula Modal -> Formula Modal -- we always use the constructor forAllB and existsB with lambda abstraction to write quantifiers
   BSub    :: [Var] -> [Expr Bool] -> Formula Modal -> Formula Modal

(∨) f g = Disj [f,g]  
(∧) f g = Conj [f,g]  
(⇒) f g = g ∨ (Neg f) 
(⇔) f g = (f ∧ g) ∨ (Neg g ∧ Neg f)
(⊕) f g = Xor [f, g] 
true = Atom (B True)
false = Atom (B False)
(⊔) p q = Nondet [p, q]  


-- | 'yvar' constructs a variable \(y_i\) for an integer i 
yvar :: Int -> Var
yvar n = BVar ("y" ++ show n)

-- | 'uvar' constructs a variable \(u_i\) for an integer i 
uvar :: Int -> Var
uvar n = BVar ("u" ++ show n)

-- | this function choose an appropriate index \(i\) to build universal quantification \(\exists y_i \ldots\)
-- from of a lambda abstracted 'FOLFormula'.
forAllB :: (Var -> Formula Modal) -> Formula Modal
forAllB f = ForAllB n body
  where body = f (uvar n)
        n    = maxUBV body + 1

-- | similar to 'forAllB' but with existential 
existsB :: (Var -> Formula Modal) -> Formula Modal
existsB f = ExistsB n body
  where body = f (yvar n)
        n    = maxYBV body + 1


-- Computes the maximum index of bound variables
maxYBV :: Formula Modal -> Int
maxYBV (Atom _  ) = 0
maxYBV (Neg f   ) = maxYBV f
maxYBV (Conj fs)  = maximum [maxYBV f | f <- fs]
maxYBV (Disj fs)  = maximum [maxYBV f | f <- fs]
maxYBV (Imp f g)  = maximum [maxYBV f, maxYBV g]
maxYBV (Equiv f g)  = maximum [maxYBV f, maxYBV g]
maxYBV (ExistsB n f) = n
maxYBV (ForAllB n f) = maxYBV f
maxYBV (K ag f) = maxYBV f 
maxYBV (Box p f) = maxYBV f 

maxUBV :: Formula Modal -> Int
maxUBV (Atom _  ) = 0
maxUBV (Neg f   ) = maxUBV f
maxUBV (Conj fs)  = maximum [maxUBV f | f <- fs]
maxUBV (Disj fs)  = maximum [maxUBV f | f <- fs]
maxUBV (Imp f g)  = maximum [maxUBV f, maxUBV g]
maxUBV (Equiv f g)  = maximum [maxUBV f, maxUBV g]
maxUBV (BSub as es g)  = maximum [maxUBV g]
maxUBV (ExistsB n f) = maxUBV f
maxUBV (ForAllB n f) = n
maxUBV (K ag f) = maxUBV f 
maxUBV (Box p f) = maxUBV f 

data Modal
type ModalFormula = Formula Modal

----------------------------
-- * Programming Language
----------------------------
data Prog
    = Test ModalFormula
    | BAssign Var BExpr 
    | Sequence [Prog]
    | Nondet [Prog]

----------------------------
-- * Strongest Precondition for comparing with IJCAI'17 paper
----------------------------
sp :: ModalFormula -> Prog -> ModalFormula
sp phi (Test beta)       = beta ∧ phi
sp phi (BAssign v e)     = existsB (\y -> Conj [Atom (BEq (BValue v) (subBExpr v (BValue y) e)), sub v (BValue y) phi])  
sp phi (Sequence [])     = phi
sp phi (Sequence [p])    = sp phi p
sp phi (Sequence (p:ps)) = sp (sp phi (Sequence ps)) p
sp phi (Nondet ps)       = Disj [sp phi p | p <- ps]

--------------------------------
-- * Variable Substitution 
-- --------------------------------

-- | sub multiple variables from a list
subList :: [(Var, Expr a)]  -> Formula b -> Formula b
subList [] f = f 
subList [(x, e)] f = sub x e f  
subList (c:cs) f = subList [c] (subList cs f)  

-- | sub is the main function, takes the variable, an expression, and the formula
-- to make the substitution
sub :: Var -> Expr a -> Formula b -> Formula b
sub x (BEq a b) f  = subBool x (BEq a b) f
sub x (B b) f      = subBool x (B b) f
sub x (BValue p) f = subBool x (BValue p) f
sub x (Xor bs) f   = subBool x (Xor bs) f
sub x (And bs) f   = subBool x (And bs) f
sub x (Or bs) f    = subBool x (Or bs) f


-- | 'subBExpr' propagates a substitution inside a boolean expression
subBExpr :: Var -> BExpr -> BExpr -> BExpr
subBExpr x e (B b)                   = B b
subBExpr x e (BValue z) | x == z     = e
subBExpr x e (BValue z) | otherwise  = BValue z
subBExpr x e (Xor bs)                = Xor [subBExpr x e b | b <- bs]
subBExpr x e (And bs)                = And [subBExpr x e b | b <- bs]
subBExpr x e (Or bs)                 = Or [subBExpr x e b | b <- bs]
subBExpr x e (BEq p q)               = BEq (subBExpr x e p) (subBExpr x e q)  

-- | 'subBool' propagates the substitution made inside a boolean expressions into a formula
subBool :: Var -> BExpr -> Formula a -> Formula a
subBool x e (Atom z)        = Atom (subBExpr x e z)
subBool x e (Neg f)         = Neg (subBool x e f)
subBool x e (Conj fs)       = Conj [subBool x e f | f <- fs]
subBool x e (Disj fs)       = Disj [subBool x e f | f <- fs]
subBool x e (Imp f g)       = Imp (subBool x e f) (subBool x e g)
subBool x e (Equiv f g)     = Equiv (subBool x e f) (subBool x e g)
subBool x e (K a f)         = K a (subBool x e f)
subBool x e (Box prog f)    = Box prog (subBool x e f)-- no used
subBool x e (ForAllB n f)   = ForAllB n (subBool x e f)  
subBool x e (ExistsB n f)   = ExistsB n (subBool x e f)  



instance Show (Formula a) where
   show f = toString f

instance Show Var where
  show (BVar a) = a

instance Show (Expr a) where
  show (BValue v) = show v
  show (BEq p q) = show p ++ "=" ++ show q
  show (B True) = "T"
  show (B False) = "F"
  show (Xor [b]) = show b
  show (Xor (b:bs)) = show b ++ "⊕" ++ show (Xor bs) 
  show (Or [b]) = show b
  show (Or (b:bs)) = show b ++ "∨" ++ show (Or bs) 
  show (And [b]) = show b
  show (And (b:bs)) = show b ++ "∧" ++ show (And bs) 


instance Show Prog where
  show (Test f) = (toString f) ++ "?"
  show (Sequence []) = ""
  show (Sequence [p]) = show p
  show (Sequence (p:ps)) = show p ++ ";" ++ show (Sequence ps) 
  show (Nondet [] ) = ""
  show (Nondet [p] ) = show p
  show (Nondet (p:ps)) = show p ++ "⊓" ++ show (Sequence ps) 
-- precedence rules to decide parentheses
precedence :: Formula t -> Int
precedence Box{} = 15
precedence ForAllB{} = 12
precedence ExistsB{} = 11
precedence Atom{} = 10
precedence Neg {} = 9
precedence K{} = 3
precedence Disj{} = 7
precedence Conj{} = 6
precedence Imp{} = 5
precedence Equiv{} = 4

-- | The main function takes any formula ('FOLFormula' or 'ModalFormula') 
-- and returns a string. The output string is formatted to be parsable by Python.
-- The plan is to use SymPy to remove redundant tautologies.
-- The elimination of the annouce logical operator exploded the size of the formula.
toString :: Formula a -> String
toString = toStringPrec 0
--

toStringPrec :: Int -> Formula a -> String
toStringPrec dCntxt f = bracket unbracketed where
    dHere       = precedence f
    recurse     = toStringPrec dHere
    unbracketed = case f of
        Atom (B t)    -> show t
        Atom (BValue (BVar b))    -> b 
        Atom (Xor (bs))  -> "Xor(" ++ separateWithCommas [recurse (Atom b) | b <- bs]  ++ ")"
        Atom e -> show e
        K ag p   -> "K_" ++ show (identity ag) ++ " " ++ recurse p
        Neg p   -> "¬" ++ recurse p
        Conj [] -> "True"
        Disj [] -> "False" 
        Conj [p] -> recurse p
        Disj [p] -> recurse p 
        Conj (p:ps) -> recurse p ++ " ∧ " ++ recurse (Conj ps) 
        Disj (p:ps) -> recurse p ++ " ∨ " ++ recurse (Disj ps) 
        Imp p q -> recurse p ++ " ⇒ " ++ recurse q
        Equiv p q -> recurse p ++ " ⇔ " ++ recurse q
        Box prog q -> "[" ++ (show prog) ++ "]" ++ recurse q
        ExistsB b f -> "∃y" ++ show b ++ "⋅"++ recurse f
        ForAllB b f -> "∀u" ++ show b ++ "⋅"++ recurse f
    bracket
        | dCntxt > dHere = \s -> "(" ++ s ++ ")"
        | otherwise      = id


separateWithCommas :: [String] -> String
separateWithCommas [b] = b
separateWithCommas (b:bs) = b ++ ", " ++ separateWithCommas bs