{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
-- |
-- Copyright   : (c) 2011,2012 Simon Meier
-- License     : GPL v3 (see LICENSE)
--
-- Portability : GHC only
--
-- Big-step proofs computing possible sources of a fact.
module Theory.Constraint.Solver.Sources (
  -- * Precomputed case distinctions

  -- ** Queries
    unsolvedChainConstraints

  -- ** Construction
  , precomputeSources
  , refineWithSourceAsms

  -- ** Application
  , solveWithSource

  -- ** Redundant cases
  , removeRedundantCases

  -- Paramters type
  , IntegerParameters(..)
  , paramOpenChainsLimit
  , paramSaturationLimit
  , showSaturationSteps

  ) where

import           Prelude                                 hiding (id, (.))
import           Safe

import qualified Data.Map                                as M
import qualified Data.Set                                as S

import           Control.Basics
import           Control.Category
import           Control.Monad.Disj
import           Control.Monad.Reader
import           Control.Monad.State                     (gets)
import           Control.Parallel.Strategies

-- uncomment for use of "EXTENSIVE_SPLIT" only
-- import           System.IO.Error
-- import           System.Environment
-- import           System.IO.Unsafe

-- import           Text.PrettyPrint.Highlight

import           Extension.Data.Label
import           Extension.Prelude

import           Theory.Constraint.Solver.Contradictions (contradictorySystem)
import           Theory.Constraint.Solver.Goals
import           Theory.Constraint.Solver.AnnotatedGoals
import           Theory.Constraint.Solver.Reduction
import           Theory.Constraint.Solver.Simplify
import           Theory.Constraint.System
import           Theory.Model

import           Control.Monad.Bind

import           Debug.Trace
import qualified GHC.Generics as G
import qualified Data.Binary  as B


-- | Parameters
data IntegerParameters = IntegerParameters
    {
      _paramOpenChainsLimit :: Integer
    , _paramSaturationLimit :: Integer
    , _showSaturationSteps  :: Bool
    } deriving( Eq, Ord, Show, G.Generic, NFData, B.Binary )
$(mkLabels [''IntegerParameters])

------------------------------------------------------------------------------
-- Precomputing case distinctions
------------------------------------------------------------------------------

-- | The number of remaining chain constraints of each case.
unsolvedChainConstraints :: Source -> [Int]
unsolvedChainConstraints =
    map (length . unsolvedChains . snd) . getDisj . get cdCases


-- Construction
---------------

-- | The initial source if the given goal is required and the
-- given source assumptions are justified.
initialSource
    :: ProofContext
    -> [LNGuarded] -- ^ Restrictions.
    -> Goal
    -> Source
initialSource ctxt restrictions goal =
    Source goal cases
  where
    polish ((name, se), _) = ([name], se)
    se0   = insertLemmas restrictions $ emptySystem RawSource $ get pcDiffContext ctxt
    cases = polish <$> runReduction instantiate ctxt se0 (avoid (goal, se0))
    instantiate = do
        insertGoal goal False
        solveGoal goal

-- | Refine a source by applying the additional proof step.
refineSource
    :: ProofContext
    -> Reduction (a, [String])  -- proof step with result and path extension
    -> Source
    -> ([a], Source)
refineSource ctxt proofStep th =
    ( map fst $ getDisj refinement
    , set cdCases newCases th )
  where
    newCases =   Disj . removeRedundantCases ctxt stableVars snd
               . map (second (modify sSubst (restrict stableVars)))
               . getDisj $ snd <$> refinement

    stableVars = frees (get cdGoal th)

    fs         = avoid th
    refinement = do
        (names, se)        <- get cdCases th
        ((x, names'), se') <- fst <$> runReduction proofStep ctxt se fs
        return (x, (combine names names', se'))

    -- Combine names such that the coerce rule is blended out.
    combine []            ns' = ns'
    combine ("coerce":ns) ns' = combine ns ns'
    combine (n       :_)  _   = [n]

-- | Solves all chain and splitting goals as well as all premise goals solvable
-- with one of the given precomputed requires source theorems, while
-- repeatedly simplifying the proof state.
--
-- Returns the names of the steps applied.
solveAllSafeGoals :: [Source] -> Integer -> Reduction [String]
solveAllSafeGoals ths' openChainsLimit =
    solve ths' [] Nothing openChainsLimit
  where
--    extensiveSplitting = unsafePerformIO $
--      (getEnv "TAMARIN_EXTENSIVE_SPLIT" >> return True) `catchIOError` \_ -> return False
    safeGoal _       _          (_,   (_, LoopBreaker)) = False
    safeGoal doSplit chainsLeft (goal, _              ) =
      case goal of
        ChainG _ _    -> if (chainsLeft > 0)
                            then True
                            else trace ("[Open Chains] Too many chain constraints, stopping precomputation. Open Chains limits (can be changed with -c=): "++ show openChainsLimit) False
        ActionG _ fa  -> not (isKUFact fa)
        -- we do not solve KD goals for Xor facts as insertAction inserts
        -- these goals directly. This prevents loops in the precomputations
        PremiseG _ fa -> not (isKUFact fa) && not (isKDXorFact fa) && not (isNoSourcesFact fa)
        DisjG _       -> doSplit
        -- Uncomment to get more extensive case splitting
        SplitG _      -> doSplit --extensiveSplitting &&
        -- SplitG _      -> False
        SubtermG _    -> doSplit

    usefulGoal (_, (_, Useful)) = True
    usefulGoal _                = False

    isKDPrem (PremiseG _ fa,_) = (isKDFact fa) && (not (isKDXorFact fa))
    isKDPrem _                 = False
    isChainPrem1 (ChainG _ (_,PremIdx 1),_) = True
    isChainPrem1 _                          = False

    solve :: [Source] -> [String] -> Maybe LNTerm -> Integer -> Reduction [String]
    solve ths caseNames lastChainTerm chainsLeft = do
        simplifySystem
        ctxt <- ask
        contradictoryIf =<< (gets (contradictorySystem ctxt))
        goals  <- gets openGoals
        chains <- gets unsolvedChains
        -- Filter out chain goals where the term in the conclusion is identical to one we just solved,
        -- as this indicates our chain can loop
        filteredGoals <- filterM  (\(g,_) -> case g of
            (ChainG c _) -> (\x -> return $ Just True /= liftM2 eqModuloFreshnessNoAC lastChainTerm x) =<< kConcTerm c
            _            -> return True) goals

        -- try to either solve a safe goal or use one of the precomputed case
        -- distinctions
        let noChainGoals = null [ () | (ChainG _ _, _) <- goals ]
            -- we perform equation splits, if there is a chain goal starting
            -- from a message variable; i.e., a chain constraint that is no
            -- open goal.
            splitAllowed    = noChainGoals && not (null chains)
            safeGoals       = fst <$> filter (safeGoal splitAllowed chainsLeft) filteredGoals
            remainingChains ((ChainG _ _):_) = chainsLeft-1
            remainingChains _                = chainsLeft
            -- we do not solve KD goals for Xor facts as insertAction inserts
            -- these goals directly. This prevents loops in the precomputations
            kdPremGoals     = fst <$> filter (\g -> isKDPrem g || isChainPrem1 g) goals
            usefulGoals     = fst <$> filter usefulGoal goals
            nextStep :: Maybe (Reduction [String], Maybe Source)
            nextStep     =
                ((\x -> (fmap return (solveGoal x), Nothing)) <$> headMay (kdPremGoals)) <|>
                ((\x -> (fmap return (solveGoal x), Nothing)) <$> headMay (safeGoals)) <|>
                (asum $ map (solveWithSourceAndReturn ctxt ths) usefulGoals)

        -- Update the last chain conclusion term if next step is a 'safe' chain goal (kdPremGoals is empty)
        lastChainTerm' <- case (kdPremGoals, safeGoals) of
            ([], ((ChainG c _):_)) -> (\t -> return $ t <|> lastChainTerm) =<< kConcTerm c
            _                      -> return lastChainTerm

        case nextStep of
          Nothing   -> return caseNames
          Just (step, Nothing) -> (\x -> solve ths (caseNames ++ x) lastChainTerm' (remainingChains safeGoals)) =<< step
          Just (step, Just usedCase) -> (\x -> solve (filterCases usedCase ths) (caseNames ++ x) lastChainTerm' (remainingChains safeGoals)) =<< step

    filterCases :: Source -> [Source] -> [Source]
    filterCases usedCase cds = filter (\x -> usedCase /= x) cds

    kConcTerm :: NodeConc -> Reduction (Maybe LNTerm)
    kConcTerm c = do
        faConc <- gets $ nodeConcFact c
        case kFactView faConc of
            Just (_,t) -> return $ Just t
            _          -> return Nothing


------------------------------------------------------------------------------
-- Redundant Case Distinctions                                              --
------------------------------------------------------------------------------

-- | Given a list of stable variables (that are referenced from outside and cannot be simply
-- renamed) and a list containing systems, this function returns a subsequence of the list
-- such that for all removed systems, there is a remaining system that is equal modulo
-- renaming of non-stable variables.
removeRedundantCases :: ProofContext -> [LVar] -> (a -> System) ->  [a] -> [a]
removeRedundantCases ctxt stableVars getSys cases0 =
    -- usually, redundant cases only occur with the multiset and bilinear pairing theories
    if enableBP msig || enableMSet msig then cases else cases0
  where
    -- decorate with index and normed version of the system
    decoratedCases = map (second addNormSys) $  zip [(0::Int)..] cases0
    -- drop cases where the normed systems coincide
    cases          =   map (fst . snd) . sortOn fst . sortednubBy (\(_,(_, x)) (_,(_, y)) -> compareSystemsUpToNewVars x y) $ decoratedCases

    addNormSys = id &&& ((modify sEqStore dropNameHintsBound) . renameDropNameHints . getSys)

    -- this is an ordering that works well in the cases we tried
    orderedVars sys =
        filter ((/= LSortNode) . lvarSort) $ map fst . sortOn snd . varOccurences $ sys

    -- rename except for stable variables, drop name hints, and import ordered vars first
    renameDropNameHints sys =
      (`evalFresh` avoid stableVars) . (`evalBindT` stableVarBindings) $ do
          _ <- renameDropNamehint (orderedVars sys)
          renameDropNamehint sys
      where
        stableVarBindings = M.fromList (map (\v -> (v, v)) stableVars)

    msig = mhMaudeSig . get pcMaudeHandle $ ctxt

------------------------------------------------------------------------------
-- Applying precomputed case distinctions
------------------------------------------------------------------------------

-- | Match a precomputed 'Source' to a goal. Returns the instantiated
-- 'Source' with the given goal if possible
matchToGoal
    :: ProofContext     -- ^ Proof context used for refining the source.
    -> Source           -- ^ Source to use.
    -> Goal             -- ^ Goal to match
    -> Maybe Source
    -- ^ An adapted version of the source with the given goal
matchToGoal ctxt th0 goalTerm =
  if not $ maybeMatcher (goalTerm, get cdGoal th0) then Nothing else
  case (goalTerm, get cdGoal th) of
    ( PremiseG      (iTerm, premIdxTerm) faTerm
     ,PremiseG pPat@(iPat,  _          ) faPat  ) ->
        case doMatch (faTerm `matchFact` faPat <> iTerm `matchLVar` iPat) of
            []      -> Nothing
            subst:_ ->
                let refine = do
                        modM sEdges (substNodePrem pPat (iPat, premIdxTerm))
                        refineSubst subst
                in Just $ snd $ refineSource ctxt refine (set cdGoal goalTerm th)

    (ActionG iTerm faTerm, ActionG iPat faPat) ->
        case doMatch (faTerm `matchFact` faPat <> iTerm `matchLVar` iPat) of
            []      -> Nothing
            subst:_ -> Just $ snd $ refineSource ctxt
                                        (refineSubst subst) (set cdGoal goalTerm th)

    -- No other matches possible, as we only precompute case distinctions for
    -- premises and KU-actions.
    _ -> Nothing
  where
    -- this code reflects the precomputed cases in 'precomputeSources'
    maybeMatcher (PremiseG _ faTerm, PremiseG _ faPat)  = factTag faTerm == factTag faPat
    maybeMatcher ( ActionG _ (Fact KUFact _ [tTerm])
                 , ActionG _ (Fact KUFact _ [tPat]))      =
        case (viewTerm tPat, viewTerm tTerm) of
            (Lit  (Var v),_) | lvarSort v == LSortFresh -> sortOfLNTerm tPat == LSortFresh
            (FApp o _, FApp o' _)                       -> o == o'
            _                                           -> True
    maybeMatcher _                                      = False

    th = (`evalFresh` avoid goalTerm) . rename $ th0

    substNodePrem from to = S.map
        (\ e@(Edge c p) -> if p == from then Edge c to else e)

    doMatch match = runReader (solveMatchLNTerm match) (get pcMaudeHandle ctxt)

    refineSubst subst = do
        void (solveSubstEqs SplitNow subst)
        void substSystem
        return ((), [])

-- | Try to solve a premise goal or 'KU' action using the first precomputed
-- case distinction with a matching premise. Also returns the used case distinction.
solveWithSourceAndReturn :: ProofContext
                         -> [Source]
                         -> Goal
                         -> Maybe (Reduction [String], Maybe Source)
solveWithSourceAndReturn hnd ths goal = asum [ applySource hnd th goal | th <- ths ]

-- | Try to solve a premise goal or 'KU' action using the first precomputed
-- source with a matching premise.
solveWithSource :: ProofContext
                -> [Source]
                -> Goal
                -> Maybe (Reduction [String])
solveWithSource hnd ths goal = fst <$> solveWithSourceAndReturn hnd ths goal

-- | Apply a precomputed source theorem to a required fact.
applySource :: ProofContext
               -> Source     -- ^ Source theorem.
               -> Goal       -- ^ Required goal
               -> Maybe (Reduction [String], Maybe Source)
applySource ctxt th0 goal = (\th -> (_applySource th, Just th0)) <$> matchToGoal ctxt th0 goal
  where
    keepVarBindings = M.fromList (map (\v -> (v, v)) (frees goal))

    _applySource :: Source -> Reduction [String]
    _applySource th = do
      markGoalAsSolved "precomputed" goal
      (names, sysTh0) <- disjunctionOfList $ getDisj $ get cdCases th
      sysTh <- evalBindT (someInst sysTh0) keepVarBindings
      conjoinSystem sysTh
      return names

-- | Saturate the sources with respect to each other such that no
-- additional splitting is introduced; i.e., only rules with a single or no
-- conclusion are used for the saturation.
saturateSources
    :: IntegerParameters -> ProofContext -> [Source] -> [Source]
saturateSources parameters ctxt thsInit  =
    go thsInit 1
  where
    go :: [Source] -> Integer -> [Source]
    go ths n
      | any or (changes `using` parList rdeepseq) && (n <= get paramSaturationLimit parameters) =
          if get showSaturationSteps parameters then
            trace ("[Saturating Sources] Step " ++ show n ++ " (Max " ++ show (get paramSaturationLimit parameters) ++ ")")
             $ go ths' (n + 1)
          else 
             go ths' (n + 1)
      | n > get paramSaturationLimit parameters =
          if get showSaturationSteps parameters then
            trace ("[Saturating Sources] Saturation aborted, more than " ++ show (get paramSaturationLimit parameters) ++
                 " iterations. (Limit can be change with -s=)") ths'
          else
            ths'
      | otherwise =
          if get showSaturationSteps parameters then
            trace "[Saturating Sources] Done" ths'
          else ths'
      where
          (changes, ths') = unzip $ map (refineSource ctxt solver) ths
          goodTh th = length (getDisj (get cdCases th)) <= 1
          solver
            = do names <- solveAllSafeGoals
                            (filter goodTh ths) (get paramOpenChainsLimit parameters)
                 return (not $ null names, names)

-- | Precompute a saturated set of case distinctions.
precomputeSources
    :: IntegerParameters
    -> ProofContext
    -> [LNGuarded]       -- ^ Restrictions.
    -> [Source]
precomputeSources parameters ctxt restrictions =
    map cleanupCaseNames (saturateSources parameters ctxt rawSources)
  where
    cleanupCaseNames = modify cdCases $ fmap $ first $
        filter (not . null)
      . map (filter (`elem` '_' : ['a'..'z'] ++ ['A'..'Z'] ++ ['0'..'9']))

    rawSources =
        (initialSource ctxt restrictions <$> (protoGoals ++ msgGoals))

    -- construct source starting from facts from non-special rules
    protoGoals = someProtoGoal <$> absProtoFacts
    msgGoals   = someKUGoal <$> absMsgFacts

    getProtoFact (Fact KUFact _ _ ) = mzero
    getProtoFact (Fact KDFact _ _ ) = mzero
    getProtoFact fa                 = return fa

    -- remove annotations to avoid precomputing the same source with multiple annotations
    absFact (Fact tag _ ts) = (tag, length ts)

    nMsgVars n = [ varTerm (LVar "t" LSortMsg i) | i <- [1..fromIntegral n] ]

    someProtoGoal :: (FactTag, Int) -> Goal
    someProtoGoal (tag, arity) =
        PremiseG (someNodeId, PremIdx 0) (Fact tag S.empty (nMsgVars arity))

    someKUGoal :: LNTerm -> Goal
    someKUGoal m = ActionG someNodeId (kuFact m)

    someNodeId = LVar "i" LSortNode 0

    -- FIXME: Also use facts from proof context.
    rules = get pcRules ctxt
    absProtoFacts = sortednub $ do
        ru <- joinAllRules rules
        fa@(tag,_) <- absFact <$> (getProtoFact =<< (get rConcs ru ++ get rPrems ru))
        -- exclude facts handled specially by the prover
        guard (not $ tag `elem` [OutFact, InFact, FreshFact])
        return fa

    absMsgFacts :: [LNTerm]
    absMsgFacts = asum $ sortednub $
      [ return $ varTerm (LVar "t" LSortFresh 1)
      -- Bilinear pairing
      , if enableBP msig then return $ fAppC EMap $ nMsgVars (2::Int) else []
      -- Natural numbers
      , if enableNat msig then
          [ fAppNoEq natOneSym []
          , fAppAC NatPlus [varTerm (LVar "t" LSortNat 1), varTerm (LVar "t" LSortNat 2)] ]
          else []
      , [ fAppNoEq o $ nMsgVars k
        | o@(_,(k,priv,_)) <- S.toList . noEqFunSyms  $ msig
        , NoEq o `S.notMember` implicitFunSig, k > 0 || priv==Private]
      ]

    msig = mhMaudeSig . get pcMaudeHandle $ ctxt

-- | Refine a set of sources by exploiting additional source
-- assumptions.
refineWithSourceAsms
    :: IntegerParameters -- ^ Parameters for openChains and Saturation limits
    -> [LNGuarded]    -- ^ Source assumptions to use.
    -> ProofContext   -- ^ Proof context to use.
    -> [Source]       -- ^ Original, raw sources.
    -> [Source]       -- ^ Manipulated, refined sources.
refineWithSourceAsms _ [] _ cases0 =
    (modify cdCases . fmap . second) (set sSourceKind RefinedSource) <$> cases0
refineWithSourceAsms parameters assumptions ctxt cases0 =
    fmap (modifySystems removeFormulas) $
    saturateSources parameters ctxt $
    modifySystems updateSystem <$> cases0
  where
    modifySystems   = modify cdCases . fmap . second
    updateSystem se =
        modify sFormulas (S.union (S.fromList assumptions)) $
        set sSourceKind RefinedSource                       $ se
    removeFormulas =
        modify sGoals (M.filterWithKey isNoDisjGoal)
      . set sFormulas S.empty
      . set sSolvedFormulas S.empty

    isNoDisjGoal (DisjG _)  _ = False
    isNoDisjGoal _          _ = True
