module Main where

import System.Exit
import System.IO
import Control.Monad
import Control.Concurrent.ParallelIO
import Control.Arrow (first)
import Data.List (partition,minimumBy,genericLength)
import Data.Ord (comparing)
import Data.Maybe
import qualified Data.ByteString.Lazy as L
import Control.Monad.Reader

import Util
import Options
import Parser
import qualified Parser.PP as PPParser
import PopulationProtocol
import Printer
import qualified Printer.DOT as DOTPrinter
import Property
import StructuralComputation
import Solver
import Solver.LayeredTermination
import Solver.StrongConsensus

writeFiles :: String -> PopulationProtocol -> [Property] -> OptIO ()
writeFiles basename pp props = do
        format <- opt outputFormat
        verbosePut 1 $ "Writing " ++ showNetName pp ++ " to " ++
                                  basename ++ " in format " ++ show format
        case format of
            OutDOT ->
                liftIO $ L.writeFile basename $ DOTPrinter.printPopulationProtocol pp

structuralAnalysis :: PopulationProtocol -> OptIO ()
structuralAnalysis pp =  do
        verbosePut 0 $ "States             : " ++ show (length (states pp))
        verbosePut 0 $ "Transitions        : " ++ show (length (transitions pp))
        verbosePut 0 $ "Initial states     : " ++ show (length (initialStates pp))
        verbosePut 0 $ "Yes states         : " ++ show (length (trueStates pp))
        verbosePut 0 $ "No states          : " ++ show (length (falseStates pp))

checkFile :: String -> OptIO PropResult
checkFile file = do
        verbosePut 0 $ "Reading \"" ++ file ++ "\""
        format <- opt inputFormat
        let parser = case format of
                             InPP -> PPParser.parseContent
        pp <- liftIO $ parseFile parser file
        props <- opt optProperties
        verbosePut 1 $ "Analyzing " ++ showNetName pp
        verbosePut 2 $
                "Number of states: " ++ show (length (states pp))
        verbosePut 2 $
                "Number of transitions: " ++ show (length (transitions pp))
        let pRowSize p = let (preP, postP) = context pp p in length preP + length postP
        let arcs = sum $ map pRowSize $ states pp
        verbosePut 2 $
                "Number of arcs: " ++ show arcs
        printStruct <- opt optPrintStructure
        when printStruct $ structuralAnalysis pp
        verbosePut 3 $ show pp
        output <- opt optOutput
        case output of
            Just outputfile ->
                writeFiles outputfile pp props
            Nothing -> return ()
        -- TODO: short-circuit? parallel?
        rs <- mapM (checkProperty pp) props
        verbosePut 0 ""
        return $ resultsAnd rs


checkProperty :: PopulationProtocol -> Property -> OptIO PropResult
checkProperty pp prop = do
        verbosePut 1 $ "\nChecking " ++ show prop
        r <- case prop of
                Correctness -> error "not yet implemented"
                LayeredTermination -> checkLayeredTermination pp
                StrongConsensus -> checkStrongConsensus pp
        verbosePut 0 $ show prop ++ " " ++ show r
        return r

printInvariant :: (Show a, Invariant a) => [a] -> OptIO ()
printInvariant inv = do
        verbosePut 0 "Invariant found"
        let invSize = map invariantSize inv
        verbosePut 2 $ "Number of atoms in invariants: " ++ show invSize ++
                " (total of " ++ show (sum invSize) ++ ")"
        mapM_ (putLine . show) inv

checkStrongConsensus :: PopulationProtocol -> OptIO PropResult
checkStrongConsensus pp = do
        r <- checkStrongConsensus' pp [] [] []
        case r of
            (Nothing, _, _, _) -> return Satisfied
            (Just _, _, _, _) -> return Unknown

checkStrongConsensus' :: PopulationProtocol ->
        [Trap] -> [Siphon] -> [StableInequality] ->
        OptIO (Maybe StrongConsensusCounterExample, [Trap], [Siphon], [StableInequality])
checkStrongConsensus' pp utraps usiphons inequalities = do
        r <- checkSat $ checkStrongConsensusSat pp utraps usiphons inequalities
        case r of
            Nothing -> return (Nothing, utraps, usiphons, inequalities)
            Just c -> do
                refine <- opt optRefinementType
                if isJust refine then
                    refineStrongConsensus pp utraps usiphons inequalities c
                else
                    return (Just c, utraps, usiphons, inequalities)

refineStrongConsensus :: PopulationProtocol ->
        [Trap] -> [Siphon] -> [StableInequality] -> StrongConsensusCounterExample ->
        OptIO (Maybe StrongConsensusCounterExample, [Trap], [Siphon], [StableInequality])
refineStrongConsensus pp utraps usiphons inequalities c@(m0, m1, m2, x1, x2) = do
        r1 <- checkSatMin $ Solver.StrongConsensus.findTrapConstraintsSat pp m0 m1 m2 x1 x2
        case r1 of
            Nothing -> do
                r2 <- checkSatMin $ Solver.StrongConsensus.findUSiphonConstraintsSat pp m0 m1 m2 x1 x2
                case r2 of
                    Nothing -> do
                        r3 <- checkSatMin $ Solver.StrongConsensus.findUTrapConstraintsSat pp m0 m1 m2 x1 x2
                        case r3 of
                            Nothing -> return (Just c, utraps, usiphons, inequalities)
                            Just utrap ->
                                checkStrongConsensus' pp (utrap:utraps) usiphons inequalities
                    Just usiphon ->
                        checkStrongConsensus' pp utraps (usiphon:usiphons) inequalities
            Just trap ->
                checkStrongConsensus' pp (trap:utraps) usiphons inequalities

checkLayeredTermination :: PopulationProtocol -> OptIO PropResult
checkLayeredTermination pp = do
        let nonTrivialTriplets = filter (not . trivialTriplet) $ generateTriplets pp
        checkLayeredTermination' pp nonTrivialTriplets 1 $ genericLength $ transitions pp

checkLayeredTermination' :: PopulationProtocol -> [Triplet] -> Integer -> Integer -> OptIO PropResult
checkLayeredTermination' pp triplets k kmax = do
        verbosePut 1 $ "Checking layered termination with at most " ++ show k ++ " layers"
        r <- checkSatMin $ checkLayeredTerminationSat pp triplets k
        case r of
            Nothing ->
                if k < kmax then
                    checkLayeredTermination' pp triplets (k + 1) kmax
                else
                    return Unknown
            Just inv -> do
                invariant <- opt optInvariant
                when invariant $ printInvariant inv
                return Satisfied

main :: IO ()
main = do
        putStrLn "Peregrine - Efficient Verification of Population Protocols\n"
        args <- parseArgs
        case args of
            Left err -> exitErrorWith err
            Right (opts, files) -> do
                when (optShowVersion opts) (exitSuccessWith "Version 0.01")
                when (optShowHelp opts) (exitSuccessWith usageInformation)
                when (null files) (exitErrorWith "No input file given")
                let opts' = opts { optProperties = reverse (optProperties opts) }
                rs <- runReaderT (mapM checkFile files) opts'
                -- TODO: short-circuit with Control.Monad.Loops? parallel
                -- execution?
                let r = resultsAnd rs
                case r of
                    Satisfied ->
                        exitSuccessWith $ "All properties " ++ show r
                    _ ->
                        exitFailureWith $ "Some properties " ++ show r

exitSuccessWith :: String -> IO ()
exitSuccessWith msg = do
        putStrLn msg
        cleanupAndExitWith ExitSuccess

exitFailureWith :: String -> IO ()
exitFailureWith msg = do
        putStrLn msg
        cleanupAndExitWith $ ExitFailure 2

exitErrorWith :: String -> IO ()
exitErrorWith msg = do
        hPutStrLn stderr msg
        cleanupAndExitWith $ ExitFailure 3

cleanupAndExitWith :: ExitCode -> IO ()
cleanupAndExitWith code = do
        stopGlobalPool
        exitWith code