module HyLo.Formula.Rewrite ( Rewr(..),
                              pnf, simplify,
                              HyLo.Formula.Rewrite.unit_tests )

where

import HyLo.Formula

import Data.List
import Control.Monad.State

import Text.Read ( Read(..) )
import qualified Text.ParserCombinators.ReadPrec as RP
import Text.ParserCombinators.ReadP    ( string )

import qualified Data.Generics.UniplateStr as Uniplate

import HyLo.Test
import HyLo.Signature.Simple ( PropSymbol )
import HyLo.Model hiding ( unit_tests )
import Test.QuickCheck

data Rewr prop = Orig prop | Rewr Int
               deriving (Eq, Ord)

instance Show prop => Show (Rewr prop) where
    showsPrec n (Orig prop) = showChar 'O' . showsPrec n prop
    showsPrec _ (Rewr i)    = showString "RW_" . shows i

instance Read prop => Read (Rewr prop) where
    readPrec = RP.choice [do 'O' <- RP.get; Orig `liftM` readPrec,
                          do _ <- RP.lift (string "RW_"); Rewr `liftM` readPrec]

wrapProps :: Formula n p r -> Formula n (Rewr p) r
wrapProps = mapSig id Orig id

data BndInfo nom = Free nom | Bound nom
                 deriving (Eq, Ord, Read, Show)

addBndInfo :: Eq n => Formula n p r -> Formula (BndInfo n) p r
addBndInfo = go Free
    where go bnd (Down i f) = Down (Bound i) $ go (record i bnd) f
          go bnd f          = onShape bnd id id (go bnd) f
          --
          record i bnd = \n -> if n == i then Bound i else bnd n

removeBndInfo :: Formula (BndInfo n) p r -> Formula n p r
removeBndInfo = mapSig unBnd id id
    where unBnd (Free  i) = i
          unBnd (Bound i) = i

type RewrFormula n p r = Formula (BndInfo n) (Rewr p) r

type Rewriter p = State [Rewr p]

runRewriter :: Rewriter p a -> a
runRewriter r = evalState r new_props
  where new_props = map Rewr [0..]

pnf :: Eq n => [Formula n p r] -> [Formula n (Rewr p) r]
pnf = runRewriter
    . fmap (map removeBndInfo . concat)
    . mapM (toPNF . addBndInfo . wrapProps . listerize . nnf)

freshProp :: Rewriter p (RewrFormula n p r)
freshProp = do (p:ps) <- get
               put ps
               return (Prop p)

-- A formula phi is in PNF (pulenta normal form) if it satisfies:
--   - no global-formula is under the scope of a modality
--     ("modality" here includes global operators, down, etc)
--   - phi is a conjunction of clauses,
--   - if <X>psi    is a subformula of phi, psi is a neg-clause
--   - if [X]psi    is a subformula of phi, psi is a clause
--   - if @s psi    is a subformula of phi, psi is a literal
--   - if dwn x psi is a subformula of phi, psi is a clause
--
--     where + a global formula is one of the form Glphi where Gl is one of:
--               * @i (but @x, with x not free, is not a global op!)
--               * [U], aka A, the universal modality
--               * <U>, aka E, the existential modality
--           + a     clause is a disjunction of literals,
--           + a neg-clause is a conjuntion of literals
--           + a literal is either:
--               * an atom
--               * a negated atom
--               * an @-formula
--               * a []-formula
--               * a <>-formula
--
-- Observation: PNF implies NNF

listerize :: Formula n p r -> Formula n p r
listerize (f :|: g) = orify  . concatMap (orList  . listerize) $ [f,g]
listerize (f :&: g) = andify . concatMap (andList . listerize) $ [f,g]
listerize f         = composeMap id listerize f

-- If f is in NNF and has its conjunctions and disjunctions "listerized", then:
--   - ^(toPNF f) is in PNF
--   - ^(toPNF f) and f are equisatisfiable
--   - all isClause (toPNF f)
toPNF :: RewrFormula n p r -> Rewriter p [RewrFormula n p r]
toPNF f
  | isDisj f                  = removeDisCon =<< mapM toPNF (orList f)
  | isConj f                  = concat `fmap` mapM toPNF (andList f)
  | Just (m,f') <- modality f = addModalOp m =<< toPNF f'
  | isAtomic f                = return [f] -- atomic or neg-atomic formula
  | otherwise                 = error $ "toPNF - unhandled case..." ++ shape f

shape :: Formula n p r -> String
shape = show . mapSig (const ()) (const ()) (const ())

-- Expects a list [l1, l2,... lk] (k > 0), representing a disjunction of
-- the conjunctions of li, where each li is a list of clauses. Observe that
-- singleton lis must be literals, and every li is non-null.
-- Returns an equisatisfiable list of clauses
removeDisCon :: [[RewrFormula n p r]] -> Rewriter p [RewrFormula n p r]
removeDisCon fs = mk_result `fmap` mapM (\_ -> freshProp) conjs
    where (lits, conjs)= partition isLit fs
            where isLit [f] = not (isDisj f)
                  isLit _   = False
          mk_result fresh_props = big_dis:lil_diss
            where big_dis  = orify $ concat (fresh_props:lits)
                  lil_diss = [(Neg p :|: f) | (p,conj) <- zip fresh_props conjs,
                                               f <- conj]

-- Expects a modality m and a list of clauses (taken as a conjunction).
-- Returns an equisatisfiable list of clauses where:
--   - m has been pushed inwards as much as possible
--   - global formulas were poped outside the m-formulas
addModalOp :: RewrModality n p r
           -> [RewrFormula n p r]
           -> Rewriter p [RewrFormula n p r]
addModalOp m fs = case (isABox m, isADia m) of
                      (True,  False) -> return $ map (addBoxToCl m) fs
                      (False, True)  -> addDiaRemoveDis m fs
                      (True,  True)  -> return $ map (addDiaToCl m) fs
                      _              -> error "addModalOp: unknown modal op"


-- Expects a modality m (a box, not self-dual) and a _clause_ c.
-- Returns a clause equivalent to prefixing m to c, without global formulas
-- under the scope of m
addBoxToCl :: RewrModality n p r -> RewrFormula n p r -> RewrFormula n p r
addBoxToCl m = orify . addBoxToF . partition isGlobalF . orList
    where addBoxToF (globals, locals) = (apply m $ orify locals) : globals

-- Expects a diamond m and a _clause_ c
-- Returns a clause equivalent to profixing m to c, without global formulas
-- under the scope of m
--
-- Observe that this function is only used on a self-dual m
addDiaToCl :: RewrModality n p r -> RewrFormula n p r -> RewrFormula n p r
addDiaToCl m = orify . addDiaToF . partition isGlobalF . orList
    where addDiaToF (globals, locals) = globals ++ map (apply m) locals

-- Expects a modality m and a list of clauses l
-- Returns a list of clauses that is equisatisfiable with m^l,
-- m is only in front of a neg-clause, no global formula is under the
-- scope of m
addDiaRemoveDis :: RewrModality n p r
                -> [RewrFormula n p r]
                -> Rewriter p [RewrFormula n p r]
addDiaRemoveDis m l = mk_res `fmap` mapM (\_ -> freshProp) disjs
    where (disjs, non_disjs) = partition isDisj l
          (globs, lits)      = partition isGlobalF non_disjs
          mk_res fresh_props = concat [
                                 [apply m $ andify (fresh_props ++ lits)],
                                 [addBoxToCl (dual m) (Neg p :|: f) |
                                                (p,f) <- zip fresh_props disjs],
                                  globs
                               ]

orList :: Formula n p r -> [Formula n p r]
orList = unfoldr next . Just
    where next (Just (a :|: b)) = Just (a, Just b)
          next (Just f)         = Just (f, Nothing)
          next Nothing          = Nothing

orify :: [Formula n p r] -> Formula n p r
orify [] = Bot
orify fs = foldr1 (:|:) fs

andList :: Formula n p r -> [Formula n p r]
andList = unfoldr next . Just
    where next (Just (a :&: b)) = Just (a, Just b)
          next (Just f)         = Just (f, Nothing)
          next Nothing          = Nothing

andify :: [Formula n p r] -> Formula n p r
andify [] = Top
andify fs = foldr1 (:&:) fs


isConj :: Formula n p r -> Bool
isConj (_ :&: _) = True
isConj _         = False

isDisj :: Formula n p r -> Bool
isDisj (_ :|: _) = True
isDisj _         = False

isAtomic :: Formula n p r -> Bool
isAtomic  Top     = True
isAtomic  Bot     = True
isAtomic (Nom  _) = True
isAtomic (Prop _) = True
isAtomic (Neg f)  = isAtomic f
isAtomic _        = False

isGlobalF :: RewrFormula n p r -> Bool
isGlobalF  = maybe False (isGlobal . fst) . modality

data Modality n p r = M_Box r | M_Dia  r
                    | M_A     | M_E
                    | M_D     | M_B
                    | M_At   n
                    | M_Down n
                    | M_Count CountOp (Where r) Int

type RewrModality n p r = Modality (BndInfo n) (Rewr p) r

isABox :: Modality n p r -> Bool
isABox M_Box {} = True
isABox M_At  {} = True
isABox M_Down{} = True
isABox M_A      = True
isABox M_D      = True
isABox _        = False

isADia :: Modality n p r -> Bool
isADia M_At  {} = True
isADia M_Down{} = True
isADia m        = not $ isABox m

isGlobal :: Modality (BndInfo n) p r -> Bool
isGlobal M_Box {}         = False
isGlobal M_Dia {}         = False
isGlobal M_Down{}         = False
isGlobal (M_At (Bound _)) = False
isGlobal M_D{}            = False
isGlobal M_B{}            = False
isGlobal _                = True

apply :: Modality n p r -> Formula n p r -> Formula n p r
apply (M_Box  r) = Box  r
apply (M_Dia  r) = Diam r
apply (M_At   i) = At   i
apply (M_Down x) = Down x
apply (M_Count c w i) = Count c w i
apply  M_A       = A
apply  M_E       = E
apply  M_D       = D
apply  M_B       = B

dual :: Modality n p r -> Modality n p r
dual (M_Box  r) = M_Dia  r
dual (M_Dia  r) = M_Box  r
dual m@M_At  {} = m
dual m@M_Down{} = m
dual  M_A       = M_E
dual  M_E       = M_A
dual  M_D       = M_B
dual  M_B       = M_D
dual (M_Count c w i) = M_Count (negCount c) w i

modality :: Formula n p r -> Maybe (Modality n p r, Formula n p r)
modality (Diam r f) = Just (M_Dia  r, f)
modality (Box  r f) = Just (M_Box  r, f)
modality (At   i f) = Just (M_At   i, f)
modality (Down x f) = Just (M_Down x, f)
modality (A f)      = Just (M_A     , f)
modality (E f)      = Just (M_E     , f)
modality (D f)      = Just (M_D     , f)
modality (B f)      = Just (M_B     , f)
modality (Count c w i f) = Just (M_Count c w i, f)
modality _          = Nothing

simplify :: (Ord n, Ord p, Ord r) => Formula n p r -> Formula n p r
simplify = Uniplate.rewrite simpl
    where simpl (Neg Top)             = Just Bot
          simpl (Neg Bot)             = Just Top
          --
          simpl (At   _ Top)          = Just Top
          simpl (At   _ Bot)          = Just Bot
          simpl (Box  _ Top)          = Just Top
          simpl (Diam _ Bot)          = Just Bot
          --
          simpl (Bot :&: _)           = Just Bot
          simpl (Top :&: f)           = Just f
          simpl (f :&: Neg g)| f == g = Just Bot
          simpl (Neg f :&: g)| f == g = Just Bot
          simpl (f :&: g)    | f == g = Just f
          --
          simpl (Top :|: _)           = Just Top
          simpl (Bot :|: f)           = Just f
          simpl (f :|: g)    | f == g = Just f
          simpl (f :|: Neg g)| f == g = Just Top
          --
          simpl (f :<-->: g) | f == g = Just Top
          simpl (Top :<-->: f)        = Just f
          simpl (Bot :<-->: f)        = Just (Neg f)
          --
          simpl (Top :-->: f)         = Just f
          simpl (Bot :-->: _)         = Just Top
          simpl (_ :-->: Top)         = Just Top
          simpl (f :-->: Bot)         = Just (Neg f)
          simpl (f :-->: g)  | f == g = Just Top
          --
          simpl f           = orient f
          --
          orient (f  :&:   g) | f > g = Just (g  :&: f)
          orient (f  :|:   g) | f > g = Just (g  :|: f)
          orient (f :<-->: g) | f > g = Just (g :<-->: f)
          orient _                  = Nothing


-- QuickCheck stuff
instance Arbitrary prop => Arbitrary (Rewr prop) where
    arbitrary = oneof [Orig `liftM` arbitrary,
                       Rewr `liftM` arbitrary]
    --
    coarbitrary (Orig prop) = variant 0 . coarbitrary prop
    coarbitrary (Rewr i)    = variant 1 . coarbitrary i

prop_read_Rewr :: Rewr PropSymbol -> Bool
prop_read_Rewr p = p == (read . show $ p)

prop_read_RewrRewr :: Rewr (Rewr PropSymbol) -> Bool
prop_read_RewrRewr p = p == (read . show $ p)

prop_simplifyIsSound :: TestFormula -> Property
prop_simplifyIsSound f = forAll (modelFor f) $ \m ->
                           forAll (worldOf m) $ \w ->
                              (m,w) |= f `iff` (m,w) |= f'
    where f' = simplify f

prop_addBndInfo :: TestFormula -> Bool
prop_addBndInfo f = null [i | Free  i <- boundVars f'] &&
                    null [i | Bound i <- freeVars  f']
    where f' = addBndInfo f

prop_removeBndInfo :: TestFormula -> Bool
prop_removeBndInfo f = f == (removeBndInfo $ addBndInfo f)

-- these properties are assumed to hold by the rewritings
precond_topBotSmall :: TestFormula -> Property
precond_topBotSmall f = (f `notElem` [Top,Bot]) ==> (f > Top && f > Bot)

unit_tests :: UnitTest
unit_tests = [("read/show Rewr",            runTest prop_read_Rewr),
              ("read/show Rewr Rewr",       runTest prop_read_RewrRewr),
              ("simplify is sound",         runTest prop_simplifyIsSound),
              ("addBndInfo binds ok",       runTest prop_addBndInfo),
              ("rmBndInfo.addBndInfo ==id", runTest prop_removeBndInfo),
              ("precond: top-bot small",    runTest precond_topBotSmall)]