{-# LANGUAGE FlexibleContexts #-}

module Solver.StrongConsensus
    (checkStrongConsensusSat,
     checkStrongConsensusCompleteSat,
     StrongConsensusCounterExample,
     RefinementObjects,
     findTrapConstraintsSat,
     findSiphonConstraintsSat,
     findUTrapConstraintsSat,
     findUSiphonConstraintsSat)
where

import Data.SBV
import qualified Data.Map as M
import Data.List ((\\),genericLength)

import Util
import PopulationProtocol
import Property
import Solver

type StrongConsensusCounterExample = (Configuration, Configuration, Configuration, FlowVector, FlowVector)
type StrongConsensusCompleteCounterExample = (Configuration, Configuration, Configuration, FlowVector, FlowVector, Configuration, Configuration, Configuration, Configuration)
type RefinementObjects = ([Trap], [Siphon])

stateEquationConstraints :: PopulationProtocol -> SIMap State -> SIMap State -> SIMap Transition -> SBool
stateEquationConstraints pp m0 m x =
            bAnd $ map checkStateEquation $ states pp
        where checkStateEquation q =
                let incoming = map addTransition $ lpre pp q
                    outgoing = map addTransition $ lpost pp q
                in  val m0 q + sum incoming - sum outgoing .== val m q
              addTransition (t,w) = literal w * val x t

nonNegativityConstraints :: (Ord a, Show a) => SIMap a -> SBool
nonNegativityConstraints m =
            bAnd $ map checkVal $ vals m
        where checkVal x = x .>= 0

terminalConstraints :: PopulationProtocol -> SIMap State -> SBool
terminalConstraints pp m =
            bAnd $ map checkTransition $ transitions pp
        where checkTransition t = bOr $ map checkState $ lpre pp t
              checkState (q,w) = val m q .<= literal (fromInteger (w - 1))

initialConfiguration :: PopulationProtocol -> SIMap State -> SBool
initialConfiguration pp m0 =
        sum (mval m0 (states pp \\ initialStates pp)) .== 0

differentConsensusConstraints :: PopulationProtocol -> SIMap State -> SIMap State -> SBool
differentConsensusConstraints pp m1 m2 =
        (sum (mval m1 (trueStates pp)) .> 0 &&& sum (mval m2 (falseStates pp)) .> 0)

unmarkedByConfiguration :: [State] -> SIMap State -> SBool
unmarkedByConfiguration r m = sum (mval m r) .== 0

markedByConfiguration :: [State] -> SIMap State -> SBool
markedByConfiguration r m = sum (mval m r) .> 0

sequenceNotIn :: [Transition] -> SIMap Transition -> SBool
sequenceNotIn u x = sum (mval x u) .== 0

sequenceIn :: [Transition] -> SIMap Transition -> SBool
sequenceIn u x = sum (mval x u) .> 0

checkUTrap :: PopulationProtocol -> SIMap State -> SIMap State -> SIMap State -> SIMap Transition -> SIMap Transition -> Trap -> SBool
checkUTrap pp m0 m1 m2 x1 x2 utrap =
            (((sequenceIn upre x1) &&& (sequenceNotIn uunmark x1)) ==> (markedByConfiguration utrap m1)) &&&
            (((sequenceIn upre x2) &&& (sequenceNotIn uunmark x2)) ==> (markedByConfiguration utrap m2))
        where upost = mpost pp utrap
              upre = mpre pp utrap
              uunmark = upost \\ upre

checkUTrapConstraints :: PopulationProtocol -> SIMap State -> SIMap State -> SIMap State -> SIMap Transition -> SIMap Transition -> [Trap] -> SBool
checkUTrapConstraints pp m0 m1 m2 x1 x2 traps =
        bAnd $ map (checkUTrap pp m0 m1 m2 x1 x2) traps

checkUSiphon :: PopulationProtocol -> SIMap State -> SIMap State -> SIMap State -> SIMap Transition -> SIMap Transition -> Siphon -> SBool
checkUSiphon pp m0 m1 m2 x1 x2 usiphon =
            (((sequenceIn upost x1) &&& (sequenceNotIn umark x1)) ==> (markedByConfiguration usiphon m0)) &&&
            (((sequenceIn upost x2) &&& (sequenceNotIn umark x2)) ==> (markedByConfiguration usiphon m0))
        where upost = mpost pp usiphon
              upre = mpre pp usiphon
              umark = upre \\ upost

checkUSiphonConstraints :: PopulationProtocol -> SIMap State -> SIMap State -> SIMap State -> SIMap Transition -> SIMap Transition -> [Siphon] -> SBool
checkUSiphonConstraints pp m0 m1 m2 x1 x2 siphons =
        bAnd $ map (checkUSiphon pp m0 m1 m2 x1 x2) siphons

checkStrongConsensus :: PopulationProtocol -> SIMap State -> SIMap State -> SIMap State -> SIMap Transition -> SIMap Transition ->
        RefinementObjects -> SBool
checkStrongConsensus pp m0 m1 m2 x1 x2 (utraps, usiphons) =
        stateEquationConstraints pp m0 m1 x1 &&&
        stateEquationConstraints pp m0 m2 x2 &&&
        initialConfiguration pp m0 &&&
        nonNegativityConstraints m0 &&&
        nonNegativityConstraints m1 &&&
        nonNegativityConstraints m2 &&&
        nonNegativityConstraints x1 &&&
        nonNegativityConstraints x2 &&&
        terminalConstraints pp m1 &&&
        terminalConstraints pp m2 &&&
        differentConsensusConstraints pp m1 m2 &&&
        checkUTrapConstraints pp m0 m1 m2 x1 x2 utraps &&&
        checkUSiphonConstraints pp m0 m1 m2 x1 x2 usiphons

checkStrongConsensusSat :: PopulationProtocol -> RefinementObjects -> ConstraintProblem Integer StrongConsensusCounterExample
checkStrongConsensusSat pp refinements =
        let m0 = makeVarMapWith ("m0'"++) $ states pp
            m1 = makeVarMapWith ("m1'"++) $ states pp
            m2 = makeVarMapWith ("m2'"++) $ states pp
            x1 = makeVarMapWith ("x1'"++) $ transitions pp
            x2 = makeVarMapWith ("x2'"++) $ transitions pp
        in  ("strong consensus", "(c0, c1, c2, x1, x2)",
             getNames m0 ++ getNames m1 ++ getNames m2 ++ getNames x1 ++ getNames x2,
             \fm -> checkStrongConsensus pp (fmap fm m0) (fmap fm m1) (fmap fm m2) (fmap fm x1) (fmap fm x2) refinements,
             \fm -> counterExampleFromAssignment (fmap fm m0) (fmap fm m1) (fmap fm m2) (fmap fm x1) (fmap fm x2))

counterExampleFromAssignment :: IMap State -> IMap State -> IMap State -> IMap Transition -> IMap Transition -> StrongConsensusCounterExample
counterExampleFromAssignment m0 m1 m2 x1 x2 =
        (makeVector m0, makeVector m1, makeVector m2, makeVector x1, makeVector x2)

-- trap and siphon refinement constraints

trapConstraint :: PopulationProtocol -> SIMap State -> Transition -> SBool
trapConstraint pp b t =
        sum (mval b $ pre pp t) .> 0 ==> sum (mval b $ post pp t) .> 0

siphonConstraint :: PopulationProtocol -> SIMap State -> Transition -> SBool
siphonConstraint pp b t =
        sum (mval b $ post pp t) .> 0 ==> sum (mval b $ pre pp t) .> 0

trapConstraints :: PopulationProtocol -> SIMap State -> SBool
trapConstraints pp b =
        bAnd $ map (trapConstraint pp b) $ transitions pp

siphonConstraints :: PopulationProtocol -> SIMap State -> SBool
siphonConstraints pp b =
        bAnd $ map (siphonConstraint pp b) $ transitions pp

uTrapConstraints :: PopulationProtocol -> FlowVector -> SIMap State -> SBool
uTrapConstraints pp x b =
        bAnd $ map (trapConstraint pp b) $ elems x

uSiphonConstraints :: PopulationProtocol -> FlowVector -> SIMap State -> SBool
uSiphonConstraints pp x b =
        bAnd $ map (siphonConstraint pp b) $ elems x

statesMarkedByConfiguration :: PopulationProtocol -> Configuration -> SIMap State -> SBool
statesMarkedByConfiguration pp m b = sum (mval b $ elems m) .> 0

statesUnmarkedByConfiguration :: PopulationProtocol -> Configuration -> SIMap State -> SBool
statesUnmarkedByConfiguration pp m b = sum (mval b $ elems m) .== 0

statesPostsetOfSequence :: PopulationProtocol -> FlowVector -> SIMap State -> SBool
statesPostsetOfSequence pp x b = sum (mval b $ mpost pp $ elems x) .> 0

statesPresetOfSequence :: PopulationProtocol -> FlowVector -> SIMap State -> SBool
statesPresetOfSequence pp x b = sum (mval b $ mpre pp $ elems x) .> 0

noOutputTransitionEnabled :: PopulationProtocol -> Configuration -> SIMap State -> SBool
noOutputTransitionEnabled pp m b =
            bAnd $ map outputTransitionNotEnabled $ transitions pp
        where
            outputTransitionNotEnabled t = outputTransitionOfB t ==> transitionNotEnabledInB t
            outputTransitionOfB t = sum [val b q | (q, w) <- lpre pp t, val m q >= w] .> 0
            transitionNotEnabledInB t = sum [val b q | (q, w) <- lpre pp t, val m q < w] .> 0

nonemptySet :: (Ord a, Show a) => SIMap a -> SBool
nonemptySet b = sum (vals b) .> 0

checkBinary :: SIMap State -> SBool
checkBinary = bAnd . map (\x -> x .== 0 ||| x .== 1) . vals

checkSizeLimit :: SIMap State -> Maybe (Int, Integer) -> SBool
checkSizeLimit _ Nothing = true
checkSizeLimit b (Just (1, curSize)) = (.< literal curSize) $ sumVal b
checkSizeLimit b (Just (2, curSize)) = (.> literal curSize) $ sumVal b
checkSizeLimit _ (Just (_, _)) = error "minimization method not supported"

minimizeMethod :: Int -> Integer -> String
minimizeMethod 1 curSize = "size smaller than " ++ show curSize
minimizeMethod 2 curSize = "size larger than " ++ show curSize
minimizeMethod _ _ = error "minimization method not supported"

findTrapConstraints :: PopulationProtocol -> StrongConsensusCounterExample -> SIMap State -> Maybe (Int, Integer) -> SBool
findTrapConstraints pp (m0, m1, m2, x1, x2) b sizeLimit =
        checkSizeLimit b sizeLimit &&&
        checkBinary b &&&
        trapConstraints pp b &&&
        (
            (statesPostsetOfSequence pp x1 b &&& statesUnmarkedByConfiguration pp m1 b) |||
            (statesPostsetOfSequence pp x2 b &&& statesUnmarkedByConfiguration pp m2 b)
        )

findTrapConstraintsSat :: PopulationProtocol -> StrongConsensusCounterExample -> MinConstraintProblem Integer Trap Integer
findTrapConstraintsSat pp c =
        let b = makeVarMap $ states pp
        in  (minimizeMethod, \sizeLimit ->
            ("trap marked by x1 or x2 and not marked in m1 or m2", "trap",
             getNames b,
             \fm -> findTrapConstraints pp c (fmap fm b) sizeLimit,
             \fm -> statesFromAssignment (fmap fm b)))

findUTrapConstraints :: PopulationProtocol -> StrongConsensusCounterExample -> SIMap State -> Maybe (Int, Integer) -> SBool
findUTrapConstraints pp (m0, m1, m2, x1, x2) b sizeLimit =
        checkSizeLimit b sizeLimit &&&
        checkBinary b &&&
        (
            (statesPostsetOfSequence pp x1 b &&& uTrapConstraints pp x1 b &&& statesUnmarkedByConfiguration pp m1 b) |||
            (statesPostsetOfSequence pp x2 b &&& uTrapConstraints pp x2 b &&& statesUnmarkedByConfiguration pp m2 b)
        )

findUTrapConstraintsSat :: PopulationProtocol -> StrongConsensusCounterExample -> MinConstraintProblem Integer Trap Integer
findUTrapConstraintsSat pp c =
        let b = makeVarMap $ states pp
        in  (minimizeMethod, \sizeLimit ->
            ("u-trap (w.r.t. x1 or x2) marked by x1 or x2 and not marked in m1 or m2", "u-trap",
             getNames b,
             \fm -> findUTrapConstraints pp c (fmap fm b) sizeLimit,
             \fm -> statesFromAssignment (fmap fm b)))

findSiphonConstraints :: PopulationProtocol -> StrongConsensusCounterExample -> SIMap State -> Maybe (Int, Integer) -> SBool
findSiphonConstraints pp (m0, m1, m2, x1, x2) b sizeLimit =
        checkSizeLimit b sizeLimit &&&
        checkBinary b &&&
        siphonConstraints pp b &&&
        statesUnmarkedByConfiguration pp m0 b &&&
        (statesPresetOfSequence pp x1 b ||| statesPresetOfSequence pp x2 b)

findSiphonConstraintsSat :: PopulationProtocol -> StrongConsensusCounterExample -> MinConstraintProblem Integer Siphon Integer
findSiphonConstraintsSat pp c =
        let b = makeVarMap $ states pp
        in  (minimizeMethod, \sizeLimit ->
            ("siphon used by x1 or x2 and unmarked in m0", "siphon",
             getNames b,
             \fm -> findSiphonConstraints pp c (fmap fm b) sizeLimit,
             \fm -> statesFromAssignment (fmap fm b)))


findUSiphonConstraints :: PopulationProtocol -> StrongConsensusCounterExample -> SIMap State -> Maybe (Int, Integer) -> SBool
findUSiphonConstraints pp (m0, m1, m2, x1, x2) b sizeLimit =
        checkSizeLimit b sizeLimit &&&
        checkBinary b &&&
        statesUnmarkedByConfiguration pp m0 b &&&
        (
            (statesPresetOfSequence pp x1 b &&& uSiphonConstraints pp x1 b) |||
            (statesPresetOfSequence pp x2 b &&& uSiphonConstraints pp x2 b)
        )

findUSiphonConstraintsSat :: PopulationProtocol -> StrongConsensusCounterExample -> MinConstraintProblem Integer Siphon Integer
findUSiphonConstraintsSat pp c =
        let b = makeVarMap $ states pp
        in  (minimizeMethod, \sizeLimit ->
            ("u-siphon (w.r.t. x1 or x2) used by x1 or x2 and unmarked in m0", "u-siphon",
             getNames b,
             \fm -> findUSiphonConstraints pp c (fmap fm b) sizeLimit,
             \fm -> statesFromAssignment (fmap fm b)))

statesFromAssignment :: IMap State -> ([State], Integer)
statesFromAssignment b = (M.keys (M.filter (> 0) b), sum (M.elems b))

-- method with all refinements directly encoded into the SMT theoryüjw

findMaximalUnmarkedTrap :: PopulationProtocol -> Integer -> SIMap Transition -> SIMap State -> SIMap State -> SBool
findMaximalUnmarkedTrap pp max x m rs =
        let stateConstraints q = unmarkedConstraints q &&& trapConstraints q
            unmarkedConstraints q = (val m q .> 0) .== (val rs q .== 0)
            trapConstraints q = (val rs q .< literal max) .== ((val rs q .== 0) ||| (successorConstraints q))
            successorConstraints q = bOr [ transitionConstraints q t | t <- post pp q ]
            transitionConstraints q t = (val x t .> 0) &&& bAnd [ val rs q' .< val rs q | q' <- post pp t ]
        in  bAnd [ stateConstraints q | q <- states pp ]

findMaximalUnmarkedSiphon :: PopulationProtocol -> Integer -> SIMap Transition -> SIMap State -> SIMap State -> SBool
findMaximalUnmarkedSiphon pp max x m s =
        findMaximalUnmarkedTrap (invertPopulationProtocol pp) max x m s

unmarkedBySequence :: PopulationProtocol -> Integer -> SIMap State -> SIMap Transition -> SBool
unmarkedBySequence pp max trap x =
            bAnd [ stateUnmarkedBySequence q | q <- states pp ]
        where stateUnmarkedBySequence q = (val trap q .== literal max) ==> sum (mval x $ pre pp q) .== 0

unusedBySequence :: PopulationProtocol -> Integer -> SIMap State -> SIMap Transition -> SBool
unusedBySequence pp max siphon x =
            bAnd [ stateUnusedBySequence q | q <- states pp ]
        where stateUnusedBySequence q = (val siphon q .== literal max) ==> sum (mval x $ post pp q) .== 0

checkBounds :: Integer -> SIMap State -> SBool
checkBounds max = bAnd . map (\x -> x .>= 0 &&& x .<= literal max) . vals

checkStrongConsensusComplete :: PopulationProtocol -> Integer -> SIMap State -> SIMap State -> SIMap State -> SIMap Transition -> SIMap Transition ->
        SIMap State -> SIMap State -> SIMap State -> SIMap State -> SBool
checkStrongConsensusComplete pp max m0 m1 m2 x1 x2 r1 r2 s1 s2 =
        stateEquationConstraints pp m0 m1 x1 &&&
        stateEquationConstraints pp m0 m2 x2 &&&
        initialConfiguration pp m0 &&&
        nonNegativityConstraints m0 &&&
        nonNegativityConstraints m1 &&&
        nonNegativityConstraints m2 &&&
        nonNegativityConstraints x1 &&&
        nonNegativityConstraints x2 &&&
        terminalConstraints pp m1 &&&
        terminalConstraints pp m2 &&&
        differentConsensusConstraints pp m1 m2 &&&
        checkBounds max r1 &&&
        checkBounds max r2 &&&
        checkBounds max s1 &&&
        checkBounds max s2 &&&
        findMaximalUnmarkedTrap pp max x1 m1 r1 &&&
        findMaximalUnmarkedTrap pp max x2 m2 r2 &&&
        findMaximalUnmarkedSiphon pp max x1 m0 s1 &&&
        findMaximalUnmarkedSiphon pp max x2 m0 s2 &&&
        unmarkedBySequence pp max r1 x1 &&&
        unmarkedBySequence pp max r2 x2 &&&
        unusedBySequence pp max s1 x1 &&&
        unusedBySequence pp max s2 x2

checkStrongConsensusCompleteSat :: PopulationProtocol -> ConstraintProblem Integer StrongConsensusCompleteCounterExample
checkStrongConsensusCompleteSat pp =
        let max = genericLength (states pp) + 1
            m0 = makeVarMapWith ("m0'"++) $ states pp
            m1 = makeVarMapWith ("m1'"++) $ states pp
            m2 = makeVarMapWith ("m2'"++) $ states pp
            x1 = makeVarMapWith ("x1'"++) $ transitions pp
            x2 = makeVarMapWith ("x2'"++) $ transitions pp
            r1 = makeVarMapWith ("r1'"++) $ states pp
            r2 = makeVarMapWith ("r2'"++) $ states pp
            s1 = makeVarMapWith ("s1'"++) $ states pp
            s2 = makeVarMapWith ("s2'"++) $ states pp
        in  ("strong consensus", "(m0, m1, m2, x1, x2, r1, r2, s1, s2)",
             concatMap getNames [m0, m1, m2, r1, r2, s1, s2] ++ concatMap getNames [x1, x2],
             \fm -> checkStrongConsensusComplete pp max (fmap fm m0) (fmap fm m1) (fmap fm m2) (fmap fm x1) (fmap fm x2) (fmap fm r1) (fmap fm r2) (fmap fm s1) (fmap fm s2),
             \fm -> completeCounterExampleFromAssignment max (fmap fm m0) (fmap fm m1) (fmap fm m2) (fmap fm x1) (fmap fm x2) (fmap fm r1) (fmap fm r2) (fmap fm s1) (fmap fm s2))

completeCounterExampleFromAssignment :: Integer -> IMap State -> IMap State -> IMap State -> IMap Transition -> IMap Transition ->
        IMap State -> IMap State -> IMap State -> IMap State -> StrongConsensusCompleteCounterExample
completeCounterExampleFromAssignment max m0 m1 m2 x1 x2 r1 r2 s1 s2 =
            (makeVector m0, makeVector m1, makeVector m2, makeVector x1, makeVector x2, makeVector r1, makeVector r2, makeVector s1, makeVector s2)
        where maximalSet q = M.keys $ M.filter (== max) q