--
-- |Providing temporal logic(s) to be as specification languages
-- and also functions to check them on a given computation tree.
--
module TL where

import AutoDriver
import Data.Maybe
import Data.List
import Control.Exception

-- |Execution sequence = the sequence subtrees along a path in a computation tree.
--
type ESeq a = [CT a]

-- |Returns the set of all execution sequences belonging to a computation tree.
--
seqOf :: CT a -> [ESeq a]
seqOf t@(BRANCH a nexts) = map (t:) . concat $ z
   where
   z = map seqOf $ nexts

-- |Representing the syntax of CTL formulas. This is not the minimalistic syntax
-- of CTL; some derived operators are included to make it easier to later on
-- normalize the negation of a CTL formula.
--
data CTL p = Now p 
           | Neg (CTL p)
           | And (CTL p) (CTL p) 
           | Or (CTL p) (CTL p) 
           | EX (CTL p) | AX (CTL p) 
           | EU (CTL p) (CTL p) 
           | AU (CTL p) (CTL p) 
           | EW (CTL p) (CTL p) 
           | AW (CTL p) (CTL p) 

           deriving (Eq, Show)
--
-- |Representing reversible predicates.
--
{-
class Reversible p where
   neg :: p -> p

instance Reversible Int where
   neg i = -i
-}

-- |To normalize a CTL formula such that negation is pushed all the way
-- towards the atomic proposition level.
--
norm (Now p) = Now p
norm (And p q) = And (norm p) (norm q)
norm (Or p q)  = Or  (norm p) (norm q)
norm (EX f)  = EX (norm f)
norm (AX f)  = AX (norm f)
norm (Neg f) = neg_ f
norm (EU f g) = EU (norm f) (norm g)
norm (AU f g) = AU (norm f) (norm g)
norm (EW f g) = EW (norm f) (norm g)
norm (AW f g) = AW (norm f) (norm g)

-- |To negate a CTL formula, and immediately normalize the result.
--
neg_ (Now p)   = Now (not . p)
neg_ (And p q) = Or (neg_ p) (neg_ q) 
neg_ (EX f)  = AX (neg_ f)
neg_ (AX f)  = EX (neg_ f)
neg_ (Neg f) = norm f
neg_ (EU f g) = AW (neg_ g) (And (neg_ f) (neg_ g))
neg_ (AU f g) = EW (neg_ g) (And (neg_ f) (neg_ g))
neg_ (EW f g) = AU (neg_ g) (And (neg_ f) (neg_ g))
neg_ (AW f g) = EU (neg_ g) (And (neg_ f) (neg_ g))

-- |Like Bool, but 'True' can be tagged with information, e.g. witness for the
-- asserted truth.
--
type BOOL u = Maybe u

tobool = isJust

-- |Representing predicates over computation tree.
--
type TPred a = CT a -> BOOL (CT a)

-- |Representating atomic/state predicates.
--
type APred a = a->Bool

-- |Representing predicates over sequence.
--
--type SPred a = ESeq a -> BOOL (CT a)

now :: APred a -> TPred a
now p (BRANCH a _) = if p a then Just (BRANCH a []) else Nothing

merge (BRANCH s1 kids1) (BRANCH s2 []) = BRANCH s1 kids1
merge (BRANCH s1 []) (BRANCH s2 kids2) = BRANCH s1 kids2
merge (BRANCH s1 kids1) (BRANCH s2 kids2) = BRANCH s1 . zipWith merge kids1 $ kids2
 
and__ p q t = and_ (p t) (q t)
   where
   and_ (Just u) (Just v) 
      | tobool . q $ u  = Just u
      | tobool . p $ v  = Just v
      | otherwise       = let
                          z = merge u v
                          pz = tobool . p $ z
                          qz = tobool . q $ z
                          in
                          if pz && qz then Just z else Just t 
                          
   and_ _ Nothing = Nothing
   and_ Nothing _ = Nothing

or__ p q t = or_ (p t) (q t)
   where
   or_ (Just u) _ = Just u
   or_ _ (Just v) = Just v
   or_ _ _        = Nothing

ex p (BRANCH a subs) = case found of
   Just t -> Just . BRANCH a $ [fromJust t] 
   _      -> Nothing
   where
   rs = map p subs
   found = find isJust rs


ax p (BRANCH a subs) = if allok then Just . BRANCH a . map fromJust $ rs else Nothing
   where
   rs = map p subs
   allok = all isJust rs 

ew p q t@(BRANCH a subs) = case q t of
   Just zq -> Just zq
   _       -> case p t of
                Nothing -> Nothing 
                Just zp -> let
                           rs = map (ew p q) subs
                           found = find isJust rs
                           zz  = fromJust . fromJust $ found
                           zz'  = BRANCH a [zz] 
                           zz'' = merge zp zz'
                           in
                           if null subs then Just zp
                           else if isNothing found then Nothing
                           else if isJust . p $ zz'  then Just zz'
                           -- this could be unsound:
                           else if isJust . p $ zz'' then Just zz''
                           else Just t

aw p q t@(BRANCH a subs) = case q t of
   Just zq -> Just zq
   _       -> case p t of
                Nothing -> Nothing 
                Just zp -> let
                           rs = map (aw p q) subs
                           ok = all isJust rs
                           zz'  = BRANCH a (map fromJust rs)
                           zz'' = merge zp zz'
                           in
                           if null subs then Just zp
                           else if not ok then Nothing
                           else if isJust . p $ zz'  then Just zz'
                           -- this could be unsound:
                           else if isJust . p $ zz'' then Just zz''
                           else Just t


eu p q t@(BRANCH a subs) = case q t of
   Just zq -> Just zq
   _       -> case p t of
                Nothing -> Nothing 
                Just zp -> let
                           rs = map (eu p q) subs
                           found = find isJust rs
                           zz  = fromJust . fromJust $ found
                           zz'  = BRANCH a [zz] 
                           zz'' = merge zp zz'
                           in
                           if isNothing found then Nothing
                           else if isJust . p $ zz'  then Just zz'
                           -- this could be unsound:
                           else if isJust . p $ zz'' then Just zz''
                           else Just t

au p q t@(BRANCH a subs) = case q t of
   Just zq -> Just zq
   _       -> case p t of
                Nothing -> Nothing 
                Just zp -> let
                           rs = map (au p q) subs
                           ok = all isJust rs
                           zz'  = BRANCH a (map fromJust rs)
                           zz'' = merge zp zz'
                           in
                           if null subs || not ok then Nothing
                           else if isJust . p $ zz'  then Just zz'
                           -- this could be unsound:
                           else if isJust . p $ zz'' then Just zz''
                           else Just t

sem (Now p) = now p
sem (And p q) = and__ (sem p) (sem q)
sem (Or p q)  = or__ (sem p) (sem q)
sem (EX p) = ex (sem p)
sem (AX p) = ax (sem p)
sem (EW p q) = ew (sem p) (sem q)
sem (AW p q) = aw (sem p) (sem q)
sem (EU p q) = eu (sem p) (sem q)
sem (AU p q) = au (sem p) (sem q)

-- |representing predicate True
tt = (\s->True)

-- |representing predicate False
ff = (\s->False)

-- |Check if a given CTL formula holds on a computation tree t. If Nothing is returned
-- then the formula holds, else Just w means that the formula does not hold, with
-- w as a prefix of t showing why the formula does not hold.
--
check ctl t = counterexample
   where
   -- Construct the negated and normalized form of the input formula:
   contra = neg_ ctl

   -- Now check the negation, to see if a counter example can be found:
   counterexample = sem contra t


checkR :: SUTstate s => Int -> CTL (s->Bool) -> SUT s -> s -> IO (BOOL (CT s))
checkR k ctl sut s0 = do {
     t <- runR k sut s0 ;
     return . check ctl $ t
   }

phix_ex :: CTL (Int->Bool)
phix_ex = EX . Now $ (==1)

phix_ax :: CTL (Int->Bool)
phix_ax = AX . Now $ (==1)

--
-- An automaton satisfying ew (<=4) F and ew (<4) (==4)
--
sutx2 :: SUT Int
sutx2 0 = return [1,5]
sutx2 1 = return [2]
sutx2 2 = return [3]
sutx2 3 = return [4]
sutx2 4 = return []
sutx2 5 = return [1]
sutx2 _ = undefined

phi_ew1,phi_ew2a,phi_ew2b,phi_ew3 :: CTL (Int->Bool)
phi_ew1 = Neg (Now (<4) `EW` Now (==4))
phi_ew2a = Neg (Now (<=4) `EW` Now ff)
phi_ew2b = Neg (Now (<4) `EW` Now ff)
phi_ew3 = Neg (Now (<4) `EW` Now (==5))

putStrln s = do { putStr s ; putStr "\n" }

test_ew = do {
   r <- checkR 2 phi_ew1 sut 0  ; check (isJust r) r ; 
   r <- checkR 2 phi_ew2a sut 0 ; check (isJust r) r ;
   r <- checkR 2 phi_ew2b sut 0 ; check (isNothing r) r ;
   r <- checkR 2 phi_ew3 sut 0  ; check (isJust r) r ;
  }
  where
  sut = traceSUT sutx2
  check b r = putStrln . show . assert b $ r

phi_eu1,phi_eu2,phi_eu3 :: CTL (Int->Bool)
phi_eu1  = Neg (Now (<4) `EU` Now (==4))
phi_eu2  = Neg (Now tt `EU` Now (==5))
phi_eu3 = Neg (Now (<4) `EU` Now (==9))

test_eu = do {
   r <- checkR 2 phi_eu1 sut 0  ; check (isJust r) r ; 
   r <- checkR 2 phi_eu2 sut 0  ; check (isJust r) r ;
   r <- checkR 2 phi_eu3 sut 0  ; check (isNothing r) r ;
  }
  where
  sut = traceSUT sutx2
  check b r = putStrln . show . assert b $ r

--
-- An automaton satisfying ew (<=5) F and ew (<5) (==5)
--
sutx3 :: SUT Int
sutx3 0 = return [1,4]
sutx3 1 = return [2]
sutx3 2 = return [3]
sutx3 3 = return [5]
sutx3 5 = return []
sutx3 4 = return [3]
sutx3 _ = undefined

phi_aw1,phi_aw2,phi_aw3 :: CTL (Int->Bool)
phi_aw1 = Neg (Now (<5) `AW` Now (==5))
phi_aw2 = Neg (Now (<=5) `AW` Now ff)
phi_aw3 = Neg (Now (<3) `AW` Now ff)

test_aw = do {
   r <- checkR 2 phi_aw1 sut 0 ; check (isJust r) r ; 
   r <- checkR 2 phi_aw2 sut 0 ; check (isJust r) r ; 
   r <- checkR 2 phi_aw3 sut 0 ; check (isNothing r) r ; 
  }
  where
  sut = traceSUT sutx3
  check b r = putStrln . show . assert b $ r

phi_au1,phi_au2,phi_au3 :: CTL (Int->Bool)
phi_au1 = Neg (Now (<5) `AU` Now (==5))
phi_au2 = Neg (Now tt `AU` Now (==3))
phi_au3 = Neg (Now (<3) `AU` Now (==5))

test_au = do {
   r <- checkR 2 phi_au1 sut 0 ; check (isJust r) r ; 
   r <- checkR 2 phi_au2 sut 0 ; check (isJust r) r ; 
   r <- checkR 2 phi_au3 sut 0 ; check (isNothing r) r ; 
  }
  where
  sut = traceSUT sutx3
  check b r = putStrln . show . assert b $ r