module DisjointTest(
                    DisjointTestState,
                    initDisjointTestState,
                    addADefToDisjointTestState,
                    checkDisjointExpr,
                    checkDisjointExprWithCtx,
                    genDisjointSet,
                    RuleDisjointTest
                   ) where

import qualified Data.Set as S
import qualified Data.Map as M
import Control.Monad(foldM {- , when -})

import Util(ordPair,uniquePairs)

import Error(ErrorHandle)
import Pretty
import PPrint(PPrint(..))

import Flags
import ASyntax
import ASyntaxUtil(AExprs(..), aAnd)
import Pragma

import VModInfo(VModInfo)
import AExpr2Util(getMethodOutputPort)
--import Debug.Trace(trace)

import qualified AExpr2STP as STP
         (SState, initSState, addADefToSState,
          checkDisjointRulePair, checkDisjointExpr)
import qualified AExpr2Yices as Yices
         (YState, initYState, addADefToYState,
          checkDisjointRulePair, checkDisjointExpr)

-- -------------------------

type RuleDisjointTest = ARuleId -> ARuleId -> Bool

-- A single data type for either of the disjoint-testing state
data DisjointTestState = DTS_Yices DSupportMap Yices.YState
                       | DTS_STP   DSupportMap STP.SState

-- -------------------------

initDisjointTestState ::
    String -> ErrorHandle -> Flags ->
    [ADef] -> [AVInst] -> [(ARuleId, [AExpr], Maybe ARuleId)] ->
    IO DisjointTestState
initDisjointTestState str errh flags ds avis rs = do
    let supportMap = buildSupportMap ds avis rs
    case (satBackend flags) of
      SAT_Yices -> do
          yices_state <- Yices.initYState str flags True ds avis rs
          return (DTS_Yices supportMap yices_state)
      SAT_STP -> do
          stp_state <- STP.initSState str flags True ds avis rs
          return (DTS_STP supportMap stp_state)


addADefToDisjointTestState :: DisjointTestState -> [ADef] ->
                              IO DisjointTestState
addADefToDisjointTestState (DTS_Yices m yices_state) ds = do
    yices_state' <- Yices.addADefToYState yices_state ds
    return (DTS_Yices m yices_state')
addADefToDisjointTestState (DTS_STP m stp_state) ds = do
    stp_state' <- STP.addADefToSState stp_state ds
    return (DTS_STP m stp_state')

-- -------------------------

checkDisjointExpr :: DisjointTestState -> AExpr -> AExpr ->
                     IO (Maybe Bool, DisjointTestState)
checkDisjointExpr (DTS_Yices m yices_state) e1 e2 = do
    (res, yices_state') <- Yices.checkDisjointExpr yices_state e1 e2
    return (res, DTS_Yices m yices_state')
checkDisjointExpr (DTS_STP m stp_state) e1 e2 = do
    (res, stp_state') <- STP.checkDisjointExpr stp_state e1 e2
    return (res, DTS_STP m stp_state')

-- When testing conditions on methods inside one rule (or two rules),
-- we also want to consider the predicates of the rules; the conditions
-- on the methods might not be disjoint, but are disjoint when the rule
-- condition is True.
--
-- XXX We assume that it's easier to test without the predicates first,
-- XXX then test again if that reports not-disjoint.
--
-- XXX For tests within the same rule, might it be cheaper to assert
-- XXX the predicate once?  Particularly if this function is being
-- XXX called many times for the same rule?
--
-- XXX This function could take the rule Ids as arguments, and look up
-- XXX the predicate from the state.
--
checkDisjointExprWithCtx ::
    DisjointTestState -> AExpr -> AExpr -> AExpr -> AExpr ->
    IO (Maybe Bool, DisjointTestState)
checkDisjointExprWithCtx dts ctx1 ctx2 e1 e2 = do
    (res, dts') <- checkDisjointExpr dts e1 e2
    case res of
      Just True -> return (res, dts')
      _ -> do -- try again but with the context included
              let e1' = aAnd ctx1 e1
                  e2' = aAnd ctx2 e2
              checkDisjointExpr dts' e1' e2'

-- -------------------------

genDisjointSet :: DisjointTestState -> [ARuleId] -> [ASchedulePragma] ->
                  IO (S.Set (ARuleId, ARuleId), DisjointTestState)
genDisjointSet dt_state ruleNames pragmas = do
  let me_state = makeMETest pragmas

      addResult :: S.Set (ARuleId, ARuleId) -> (ARuleId, ARuleId) -> Maybe Bool -> S.Set (ARuleId, ARuleId)
      addResult rset (r1,r2) (Just True) = S.insert (r1,r2) (S.insert (r2,r1) rset)
      addResult rset _ _                 = rset

      foldFn (rset, st) p =
          case (checkMERulePair me_state p) of
            res@(Just _) -> let rset' = addResult rset p res
                            in  return (rset', st)
            _            -> do (res, st') <- checkDisjointRulePairTop st p
                               let rset' = addResult rset p res
                               return (rset', st')

  (res, dt_state') <- foldM foldFn (S.empty, dt_state) (uniquePairs ruleNames)

  return (res, dt_state')


supportIntersects :: DSupportMap -> (ARuleId, ARuleId) -> Bool
supportIntersects m (r1, r2) = S.null $ S.intersection s1 s2
    where s1 = lookupSupport m r1
          s2 = lookupSupport m r2


checkDisjointRulePairTop :: DisjointTestState -> (ARuleId, ARuleId) ->
                         IO (Maybe Bool, DisjointTestState)
checkDisjointRulePairTop s p = do
  let dis = supportIntersects (getSupportMap s) p
  if (dis) then return (Just False,s)
    else do checkDisjointRulePair s p
  -- Self check code.....
  -- (res, s') <-checkDisjointRulePair s p
  -- when (dis && (res /= Just False)) $ do
  --   putStrLn $ "DIS " ++ ppReadable p ++ " --> " ++ ppReadable (res, dis)
  --   putStrLn $ "XXXXXX " ++ ppReadable (getSupportMap s)
  --   putStrLn $ "XXXXXX " ++ ppReadable res
  --   fail "bad disjoint map"
  -- return (res, s')


checkDisjointRulePair :: DisjointTestState -> (ARuleId, ARuleId) ->
                         IO (Maybe Bool, DisjointTestState)
checkDisjointRulePair s@(DTS_Yices m yices_state) p = do
    (res, yices_state') <- Yices.checkDisjointRulePair yices_state p
    return (res, DTS_Yices m yices_state')
checkDisjointRulePair s@(DTS_STP m stp_state) p = do
    (res, stp_state') <- STP.checkDisjointRulePair stp_state p
    return (res, DTS_STP m stp_state')

-- -------------------------

type MEMap  = M.Map ARuleId [Integer] -- map rid to a list of me sets
type METest = M.Map (Integer,Integer) (Maybe Bool)


checkMERulePair :: (MEMap, METest) -> (ARuleId,ARuleId) -> Maybe Bool
checkMERulePair (me_map, me_test) (rule1,rule2) =
  case (M.lookup rule1 me_map, M.lookup rule2 me_map) of
          (Just nums1, Just nums2) -> checkMEPairs me_test nums1 nums2
          _                        -> Nothing

checkMEPairs :: METest -> [Integer] -> [Integer] -> Maybe Bool
checkMEPairs me_test [] _  = Nothing
checkMEPairs me_test _  [] = Nothing
checkMEPairs me_test (num1:rest) nums2 =
      let fn [] = Nothing
          fn (num2:rest) = case (M.lookup (ordPair (num1, num2)) me_test) of
                             (Just res) -> res
                             _          -> fn rest
      in  case (fn nums2) of
            (Just y) -> (Just y)
            _        -> checkMEPairs me_test rest nums2


makeMETest :: [ASchedulePragma] -> (MEMap, METest)
makeMETest spms = fst $ foldl addMESPM ((M.empty, M.empty), 0) spms

addMESPM :: ((MEMap, METest), Integer) -> ASchedulePragma ->
            ((MEMap, METest), Integer)
addMESPM s (SPMutuallyExclusive ids) = addMEIds s ids
addMESPM s _ = s

addMEIds :: ((MEMap, METest), Integer) -> [[ARuleId]] ->
            ((MEMap, METest), Integer)
addMEIds s idss =
  let -- XXX this is also in ASchedule
      mkMEPairs [] = []
      mkMEPairs ([]:rest) = mkMEPairs rest
      mkMEPairs (ids:rest) = [(ids, r) | r <- rest] ++ (mkMEPairs rest)

      addMEPair (ids0, ids1) ((me_map, me_test), num) =
          let me_map1 = foldl (addMEId num) me_map ids0
              me_map2 = foldl (addMEId (num+1)) me_map1 ids1
              me_test1 = M.insert (num, num+1) (Just True) me_test
          in  ((me_map2, me_test1), num+2)
  in
      foldr addMEPair s (mkMEPairs idss)

addMEId :: Integer -> MEMap -> ARuleId -> MEMap
addMEId num me_map id = do
  case (M.lookup id me_map) of
      Just nums -> M.insert id (num:nums) me_map
      Nothing   -> M.insert id [num]      me_map

-- -------------------------
-- Disjoint testing can be avoided if there is no overlap i the logic cones of each rule,
-- that is f(as) == 1 and g(bs) == 1 can be satisfied if set as /= set bs
-- we build a map containing the support set (as, bs) for each Aid,  rule
type DSupportMap =  M.Map AId (S.Set ASupport)
data ASupport = DMethod AId AId -- state port
                | DLeaf AId
                | DDef ADef AId
                | DClkGate AId AId
                | DTask AId
                deriving (Show, Eq, Ord)


isLeafDef :: ASupport -> Bool
isLeafDef (DMethod _ _)  = True
isLeafDef (DLeaf _)      = True
isLeafDef (DClkGate _ _) = True
isLeafDef (DTask _)      = True
isLeafDef (DDef _ _)     = False

lookupSupport :: DSupportMap -> AId -> (S.Set ASupport)
lookupSupport m d = M.findWithDefault err d m
    where err = error $ "DisjointTests findSupport, bad lookup: " ++ show d
                ++ "\n" ++ show m

buildSupportMap :: [ADef] -> [AVInst] -> [(ARuleId, [AExpr], Maybe ARuleId)] -> DSupportMap
buildSupportMap adefs avis rs = --trace ("XXX support map:" ++ ppReadable res) $
                                res
  where
    res = foldl generator M.empty [(id,es) | (id,es,_) <- rs]
    --
    idToDef:: M.Map AId ADef
    idToDef     = M.fromList [ (id, def) | def@(ADef id _ _ _) <- adefs]
    err i = error $ "DisjointTest::buildSupportMap Invalid lookup in idtodefmap: " ++ show i ++ "\n" ++ show idToDef
    --
    portMap :: M.Map AId VModInfo
    portMap = M.fromList [(avi_vname avi, avi_vmi avi) | avi <- avis]
    --
    generator :: (AExprs a) => DSupportMap -> (AId,a) -> DSupportMap
    generator m (id, es) | M.member id m = m
                         | otherwise = m''
      where
        fanins = findAExprs findSupport es
        faninIds = [ id | (DDef def id) <- fanins]
        m' = foldl generator m [ (id,def) | (DDef def id) <- fanins]
        --
        localSupports = S.fromList $ filter isLeafDef fanins
        childSupports = S.unions $ map (lookupSupport m') faninIds
        m'' = M.insert id (localSupports `S.union` childSupports) m'
    --
    findSupport                                          :: AExpr -> [ASupport]
    findSupport e@(ASDef _ i)                            = [DDef def i]
        where def = M.findWithDefault (err i) i idToDef
    findSupport e@(APrim { ae_args = es})                = findAExprs findSupport es
    findSupport e@(AMethCall {ae_args = es})             = findAExprs findSupport es ++ [DMethod (ae_objid e) vlogport]
      where vlogport = getMethodOutputPort portMap (ae_objid e) (ameth_id e)
    findSupport e@(AMethValue {})                        = [DMethod (ae_objid e) (ameth_id e)]
    findSupport e@(ANoInlineFunCall{ ae_args = es})      = findAExprs findSupport es
    findSupport e@(ATaskValue {ae_objid=id})             = [DTask id]
    findSupport e@(ASPort {ae_objid = id})               = [DLeaf id]
    findSupport e@(ASParam {ae_objid = id})              = [DLeaf id]
    findSupport e@(ASInout {ae_expi = aio})              = findAExprs findSupport wires
        where wires = ainout_wire aio
    findSupport e@(AMGate {ae_objid = id, ae_clkid=clk}) = [DClkGate id clk]
    findSupport _                                        = []

-- -------------------------
getSupportMap :: DisjointTestState -> DSupportMap
getSupportMap (DTS_Yices m _) = m
getSupportMap (DTS_STP m _)   = m


instance PPrint ASupport where
    pPrint d _ (DMethod i m) = parens $ text "DMethod" <+> pPrint d 0 (i,m)
    pPrint d _ (DLeaf   i)   = parens $ text "DLeaf" <+> pPrint d 0 i
    pPrint d _ (DDef _ i)    = parens $ text "DDef" <+> pPrint d 0 i
    pPrint d _ (DClkGate m c) = parens $ text "DClkGate" <+> pPrint d 0 (m,c)
    pPrint d _ (DTask i)    = parens $ text "DTask" <+> pPrint d 0 i