{-# LANGUAGE CPP, DeriveDataTypeable, PatternGuards #-}
module Language.Java.Paragon.TypeCheck.Policy 
    (
     bottom, top, thisP, join, joinWThis, meet, {-recmeet,-} isTop, isBottom, includesThis,
     TcPolicy(..), PrgPolicy(..),{-TcPolicyRec(..),-} TcClause(..), TcAtom(..), TcActor(..), TcMetaVar(..),
     {- toPolicyLit, -} {-toRecPolicy,-}
     lockToAtom, 
     {-zonkPolicy,-} substPolicy, firstRigid, substThis,
     flowAtomString,

     ActorPolicy, AtomPolicy
    ) where

import Language.Java.Paragon.Syntax
import Language.Java.Paragon.Pretty
import Language.Java.Paragon.Interaction

import Language.Java.Paragon.TypeCheck.Actors
import Language.Java.Paragon.TypeCheck.Locks

import Data.IORef
import Data.List ( (\\) {-, groupBy, nub-}  )
import Data.Maybe 

#ifdef BASE4
import Data.Data
#else
import Data.Generics (Data(..),Typeable(..))
#endif

flowAtomString :: String
flowAtomString = "$FLOW$"

policyModule :: String
policyModule = typeCheckerBase ++ ".Policy"

---------------------------------------------------------------
-- Basic policies

data TcPolicy a = RealPolicy (PrgPolicy a) 
                | Join (TcPolicy a) (TcPolicy a)
                | VarPolicy (TcMetaVar a)
  deriving (Eq, Ord, Show, Data, Typeable)

data PrgPolicy a = TcPolicy [TcClause a]
                 | TcRigidVar (Ident ())
                 | TcThis
                 | TcJoin (PrgPolicy a) (PrgPolicy a)
                 | TcMeet (PrgPolicy a) (PrgPolicy a)
  deriving (Eq,Ord,Show,Data,Typeable)
  

{-
data VarPolicy a = VPolicy [TcClause a]
                 | VRigidVar (Ident ())
                 | VThis
                 | VJoin (VarPolicy a) (VarPolicy a)
                 | VMeet (VarPolicy a) (VarPolicy a)
                 | VMetaVar a
  deriving (Eq,Ord,Show,Data,Typeable)
-}

type ActorPolicy = TcPolicy TcActor

type AtomPolicy = TcPolicy TcAtom

data TcMetaVar a = TcMetaVar Int (IORef (Maybe (TcPolicy a)))
--  deriving (Data, Typeable)

instance Data (TcMetaVar a) where
    gunfold    = panic (policyModule ++ ": instance Data TcMetaVar (gunfold)")
                 "Meta variables should never be reified"
    toConstr   = panic (policyModule ++ ": instance Data TcMetaVar (toConstr)")
                 "Meta variables should never be reified"
    dataTypeOf = panic (policyModule ++ ": instance Data TcMetaVar (dataTypeOf)")
                 "Meta variables should never be reified"

instance Typeable (TcMetaVar a) where
    typeOf = panic (policyModule ++ ": instance Typeable TcMetaVar (typeOf)")
             "Meta variables should never be reified"

instance Eq (TcMetaVar a) where
  (TcMetaVar i1 _) == (TcMetaVar i2 _) = i1 == i2

instance Ord (TcMetaVar a) where
  compare (TcMetaVar i1 _) (TcMetaVar i2 _) = compare i1 i2

instance Show (TcMetaVar a) where
  show (TcMetaVar i _) = "$$" ++ show i

data TcClause a = TcClause a [TcAtom]
  deriving (Eq,Ord,Show,Data,Typeable)

data TcAtom = TcAtom (Name ()) [TcActor]
  deriving (Eq,Ord,Show,Data,Typeable)

data TcActor = TcActor ActorId
             | TcVar (Ident ())
  deriving (Eq,Ord,Show,Data,Typeable)

---------------------------------------------
-- Pretty printing

instance Pretty a => Pretty (PrgPolicy a) where
  pretty (TcPolicy []) = braces $ char ':'
  pretty (TcPolicy cs) = braces $ hcat (punctuate (char ';') $ map pretty cs)
--  pretty (TcPolicyVar mv) = pretty mv
  pretty (TcRigidVar i) = text "policyof" <> parens (pretty i)
  pretty (TcThis)       = text "policyof" <> parens (text "this")
  pretty (TcJoin p1 p2) = pretty p1 <+> char '*' <+> pretty p2
  pretty (TcMeet p1 p2) = pretty p1 <+> char '+' <+> pretty p2

instance Pretty a => Pretty (TcPolicy a) where
    pretty (RealPolicy p) = pretty p
    pretty (Join p1 p2) = pretty p1 <+> char '*' <+> pretty p2
    pretty (VarPolicy mv) = pretty mv

instance Pretty (TcMetaVar a) where
  pretty (TcMetaVar i _) = text ("$$" ++ show i)

instance Pretty a => Pretty (TcClause a) where
  pretty (TcClause h b) = pretty h <> char ':' <+> hcat (punctuate (char ',') $ map pretty b)

instance Pretty TcAtom where
  pretty (TcAtom n as) = pretty n <> opt (not $ null as) (parens (hcat (punctuate (char ',') $ map pretty as)))

instance Pretty TcActor where
  pretty (TcActor aid) = pretty aid
  pretty (TcVar i) = char '\'' <> pretty i

mkSimpleLName :: Ident () -> Name ()
mkSimpleLName i = Name () LName Nothing i

--------------------------------------
-- Converting policies to recursive --
-- ones, i.e. with the head flow    --
--------------------------------------


{-TO FIX
toRecPolicy :: TcPolicy -> TcPolicyRec
toRecPolicy (TcPolicy cs)  = TcPolicyRec $ map toRecClause cs
  where toRecClause (TcClause act ats) = TcClause (TcAtom (mkSimpleLName (Ident () flowAtomString)) [act]) ats
toRecPolicy (TcJoin p1 p2) = TcJoinRec (toRecPolicy p1) (toRecPolicy p2)
toRecPolicy (TcMeet p1 p2) = TcMeetRec (toRecPolicy p1) (toRecPolicy p2)
toRecPolicy _ = panic (policyModule ++ "toRecPolicy")
                "Cannot convert non-literal policy to a recursive one"
-}



{-----------------------------------
-- Turning type-level policies   --
-- back to their source reps     --
-----------------------------------

toActor :: TcActor -> Actor
toActor (TcVar i) = Var i
toActor (TcActor aid) = Actor $ ActorName $ Name $ [Ident $ "a" ++ show (getId aid)] -- reprName aid

toActorName :: ActorId -> ActorName
toActorName (Fresh k) = ActorName $ Name $ [Ident $ "a" ++ show k]
toActorName (Alias k) = ActorName $ Name $ [Ident $ "a" ++ show k]

toAtom :: TcAtom -> Atom
toAtom (TcAtom n acts) = Atom n $ map toActor acts

toClause :: TcClause TcActor -> Clause ActorName
toClause (TcClause h ats) = Clause (toActor h) $ map toAtom ats

toPolicyLit :: TcPolicy -> PolicyExp
toPolicyLit (TcPolicy cs) = PolicyLit $ map toClause cs

----------------------------------}
-- Basic policies

class IsPolicy p where
  toPolicy :: PrgPolicy a -> p a
  fromPolicy :: p a -> Maybe (PrgPolicy a)
  includesThis :: p a -> Bool

instance IsPolicy PrgPolicy where
  toPolicy = id
  fromPolicy = Just
  
  includesThis TcThis = True
  includesThis (TcJoin p q) = any includesThis [p,q]
  includesThis (TcMeet p q) = any includesThis [p,q]
  includesThis _ = False

instance IsPolicy TcPolicy where
  toPolicy = RealPolicy
  fromPolicy (RealPolicy p) = Just p
  fromPolicy _ = Nothing

  includesThis (Join p q) = any includesThis [p,q]
  includesThis (RealPolicy p) = includesThis p
  includesThis _ = False

bottom, top, thisP :: IsPolicy pol => pol TcActor
bottom = toPolicy $ TcPolicy [TcClause (TcVar $ Ident () "x") []]
top = toPolicy $ TcPolicy []
thisP = toPolicy TcThis

isTop, isBottom :: IsPolicy pol => pol TcActor -> Bool
isTop pol | Just (TcPolicy []) <- fromPolicy pol = True
isTop _ = False

isBottom pol | Just (TcPolicy cs) <- fromPolicy pol = any isBottomC cs
  where isBottomC (TcClause (TcVar _) []) = True
        isBottomC _ = False
isBottom _ = False

----------------------------------
-- meet and join

joinWThis :: (TcPolicy TcActor) -> (TcPolicy TcActor) -> (TcPolicy TcActor)
joinWThis (RealPolicy p) (RealPolicy q) = RealPolicy $ lubWThis p q 
joinWThis p q = Join p q

lubWThis :: (PrgPolicy TcActor) -> (PrgPolicy TcActor) -> (PrgPolicy TcActor)
lubWThis TcThis q = q
lubWThis p TcThis = p
lubWThis p q = lub p q

lub :: (PrgPolicy TcActor) -> (PrgPolicy TcActor) -> (PrgPolicy TcActor)
lub (TcPolicy cs1) (TcPolicy cs2) =
      let sameFixedCs = [ TcClause (TcActor aid) (as ++ substAll senv bs) 
                              | TcClause (TcActor aid) as <- cs1,
                                TcClause (TcActor aid') bs <- cs2,
                                aid == aid',
                                let senv = zip (fuv bs) (map TcVar $ allBinders \\ fuv as)]
          bothVarCs   = [ TcClause a1 (as ++ substAll senv bs) 
                              | TcClause a1@(TcVar _) as <- cs1,
                                TcClause (TcVar i2) bs <- cs2,
                                let senv = (i2,a1): zip (fuv bs) (map TcVar $ allBinders \\ fuv as) ]
          fstVarCs    = [ TcClause a1 (as ++ substAll senv bs)
                              | TcClause a1@(TcActor _) as <- cs1,
                                TcClause (TcVar i2) bs <- cs2,
                                let senv = (i2,a1): zip (fuv bs) (map TcVar $ allBinders \\ fuv as) ]
          sndVarCs    = [ TcClause a2 (bs ++ substAll senv as)
                              | TcClause (TcVar i1) as <- cs1,
                                TcClause a2@(TcActor _) bs <- cs2,
                                let senv = (i1,a2): zip (fuv as) (map TcVar $ allBinders \\ fuv bs) ]
      in {- return $ -} TcPolicy $ sameFixedCs ++ bothVarCs ++ fstVarCs ++ sndVarCs

lub p1 p2 = TcJoin p1 p2

join :: (TcPolicy TcActor) -> (TcPolicy TcActor) -> (TcPolicy TcActor) 
join (RealPolicy p) (RealPolicy q) = RealPolicy $ p `lub` q
join p q = Join p q
                          

type Subst = [(Ident (), TcActor)]

--subst :: Ident () -> TcActor -> [TcAtom] -> [TcAtom]
--subst i1 i2 = substAll [(i1,i2)]

substAll :: Subst -> [TcAtom] -> [TcAtom]
substAll s = map (substAtom s)
    where substAtom env (TcAtom n acts) = TcAtom n (map (substA env) acts)

substA :: Subst -> TcActor -> TcActor
substA env (TcVar i) = case lookup i env of
                         Nothing -> TcVar i
                         Just a -> a
substA _ act = act

fuv :: [TcAtom] -> [Ident ()]
fuv = concatMap fuv'
    where fuv' (TcAtom _ as) = concatMap fuvA as

fuvA :: TcActor -> [Ident ()]
fuvA (TcActor _) = []
fuvA (TcVar i) = [i]

allBinders :: [Ident ()]
allBinders = map (Ident ()) $ 
              [ show c | c <- ['a' .. 'z']] ++
               [ c : show (i :: Int) | c <- ['a' .. 'z'], i <- [0..]]


glb :: (PrgPolicy TcActor) -> (PrgPolicy TcActor) -> (PrgPolicy TcActor)
-- This one could be smartified
glb (TcPolicy as) (TcPolicy bs) = TcPolicy $ as ++ bs
glb p1 p2 = TcMeet p1 p2

meet :: (TcPolicy TcActor) -> (TcPolicy TcActor) -> (TcPolicy TcActor)
meet (RealPolicy p) (RealPolicy q) = RealPolicy $ glb p q
meet _ _ = panic (policyModule ++ ".meet") 
           "meet used on non-programmable policies"

{-TO FIX
recmeet :: TcPolicyRec -> TcPolicyRec -> TcPolicyRec
recmeet (TcPolicyRec cs1) (TcPolicyRec cs2) = TcPolicyRec $ cs1 ++ cs2
recmeet p1                p2                = TcMeetRec p1 p2
-}

----------------------------------------------
{-- Specialisation

specialise :: [TcLock] -> PrgPolicy TcActor -> PrgPolicy TcActor
specialise ls (TcPolicy cs) = TcPolicy $ concatMap (specClause ls) cs
specialise _ p = p

specClause :: [TcLock] -> TcClause TcActor -> [TcClause TcActor]
specClause ls (TcClause h atoms) = 
    -- Step 1: generate possible substitutions from each atom
    let substss = map (genSubst ls) atoms -- :: [[Subst]]
        substs  = joinSubsts $ map ([]:) substss
    in remSubsumed [] $
           [ TcClause (substA sub h) (substAll sub atoms \\ map lockToAtom ls) 
             | sub <- substs ]

  where genSubst :: [TcLock] -> TcAtom -> [Subst]
        genSubst locks (TcAtom n acts) =
            [ subst' | (TcLock ln aids) <- locks,
                       n == ln,
                       let mS = fmap concat $ sequence $ map (uncurry matchesActor) (zip acts aids),
                       isJust mS,
                       let sub = fromJust mS
                           mSubst = toConsistent sub,
                       isJust mSubst, -- is this redundant?
                       let Just subst' = mSubst ]

        joinSubsts :: [[Subst]] -> [Subst]
        joinSubsts sss = 
            let allS = map concat $ cartesian sss -- [Subst]
                conS = catMaybes $ map toConsistent allS      -- [Subst]
            in conS

        remSubsumed :: [TcClause TcActor] -> [TcClause TcActor] -> [TcClause TcActor]
        remSubsumed acc [] = acc
        remSubsumed acc (x:xs) = if any (subsumes x) (acc ++ xs)
                                  then remSubsumed acc xs
                                  else remSubsumed (x:acc) xs

        subsumes :: TcClause TcActor -> TcClause TcActor -> Bool
        subsumes (TcClause h1 as1) (TcClause h2 as2) =
            h1 == h2 && null (as2 \\ as1)

        cartesian :: [[Subst]] -> [[Subst]]
        cartesian [] = [[]]
        cartesian (xs:xss) = [ y:ys | y <- xs, ys <- cartesian xss ]

matchesActor :: TcActor -> ActorId -> Maybe [(Ident (), TcActor)]
matchesActor (TcActor aid1) aid2 | aid1 == aid2 = Just []
matchesActor (TcVar i) aid = Just [(i, TcActor aid)]
matchesActor _ _ = Nothing

toConsistent :: Subst -> Maybe Subst
toConsistent sub = 
    let bndrs = groupBy (\a b -> fst a == fst b) sub
    in  sequence $ map checkConsistent bndrs
        where checkConsistent :: Subst -> Maybe (Ident (), TcActor)
              checkConsistent s = case nub s of
                                    [x] -> Just x
                                    _   -> Nothing

-}
lockToAtom :: TcLock -> TcAtom
lockToAtom (TcLock n aids) = TcAtom n $ map TcActor aids
lockToAtom (TcLockVar i) = TcAtom (mkSimpleLName i) []

{-
---------------------------
-- Zonking

zonkPolicy :: TcPolicy -> IO TcPolicy
zonkPolicy p@(TcPolicyVar (TcMetaVar _i ref)) = do
  mPol <- readIORef ref
  case mPol of
    Just pol -> do pol' <- zonkPolicy pol
                   writeIORef ref (Just pol') -- short-circuit
                   return pol'
    Nothing -> return p

zonkPolicy (TcJoin p1 p2) = do
  p1' <- zonkPolicy p1
  p2' <- zonkPolicy p2
  return $ p1' `join` p2'

zonkPolicy (TcMeet p1 p2) = do
  p1' <- zonkPolicy p1
  p2' <- zonkPolicy p2
  return $ p1' `meet` p2'

zonkPolicy p = return p
-}
-------------------------------
-- Policy substitution

-- Invariant: policies are always closed by the env
substPolicy :: [(Ident (), PrgPolicy TcActor)] -> PrgPolicy TcActor -> PrgPolicy TcActor
substPolicy env (TcRigidVar i) 
    | Just q <- lookup i env = q
substPolicy env (TcJoin p1 p2) =
    let pol1 = substPolicy env p1
        pol2 = substPolicy env p2
    in pol1 `lub` pol2
substPolicy env (TcMeet p1 p2) =
    let pol1 = substPolicy env p1
        pol2 = substPolicy env p2
    in pol1 `glb` pol2
substPolicy _ p = p

firstRigid :: PrgPolicy a -> Maybe (Ident ())
firstRigid (TcRigidVar i) = Just i
firstRigid (TcJoin p q) = listToMaybe $ catMaybes $ map firstRigid [p,q]
firstRigid (TcMeet p q) = listToMaybe $ catMaybes $ map firstRigid [p,q]
firstRigid _ = Nothing

substThis :: PrgPolicy TcActor -> PrgPolicy TcActor -> PrgPolicy TcActor
substThis p TcThis = p
substThis p (TcJoin p1 p2) =
    let pol1 = substThis p p1
        pol2 = substThis p p2
    in pol1 `lub` pol2
substThis p (TcMeet p1 p2) =
    let pol1 = substThis p p1
        pol2 = substThis p p2
    in pol1 `glb` pol2
substThis _ p = p

{-
mkPolicySubst :: [TcPolicy] -> [TcPolicy] -> [(Ident (), TcPolicy)]
mkPolicySubst pPars pArgs = catMaybes $ map aux $ zip pPars pArgs
  where aux (TcRigidVar i, p) = Just (i,p)
        aux _ = Nothing
-}