{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}

module Parser.PP
    (parseContent)
where

import Data.Aeson
import Data.Aeson.TH (deriveJSON, defaultOptions)
import Control.Applicative ((<*),(*>),(<$>))
import Data.Functor.Identity
import Text.Parsec
import qualified Data.Set as S
import qualified Data.ByteString.Lazy.Char8 as BS
import qualified Data.Text as T
import Text.Parsec.Expr
import Text.Parsec.Language (LanguageDef, emptyDef)
import qualified Text.Parsec.Token as Token
import Parser
import PopulationProtocol (PopulationProtocol,makePopulationProtocolFromStrings,State(..),Transition(..))
import Property

languageDef :: LanguageDef ()
languageDef =
        emptyDef {
                 Token.commentStart    = "/*",
                 Token.commentEnd      = "*/",
                 Token.commentLine     = "//",
                 Token.identStart      = letter <|> char '_',
                 Token.identLetter     = alphaNum <|> char '_',
                 Token.reservedNames   = ["true", "false", "EXISTS", "FORALL"],
                 Token.reservedOpNames = ["->", "<", "<=", "=", "!=", ">=", ">",
                                          "+", "-", "*", "&&", "||", "!", ":"]
                 }
 
lexer :: Token.TokenParser ()
lexer = Token.makeTokenParser languageDef

identifier :: Parser String
identifier = Token.identifier lexer -- parses an identifier
stringLiteral :: Parser String
stringLiteral = Token.stringLiteral lexer -- parses a string literal
reserved :: String -> Parser ()
reserved   = Token.reserved   lexer -- parses a reserved name
reservedOp :: String -> Parser ()
reservedOp = Token.reservedOp lexer -- parses an operator
brackets :: Parser a -> Parser a
brackets   = Token.brackets   lexer -- parses p surrounded by brackets
braces :: Parser a -> Parser a
braces     = Token.braces     lexer -- parses p surrounded by braces
parens :: Parser a -> Parser a
parens     = Token.parens     lexer -- parses p surrounded by parenthesis
natural :: Parser Integer
natural    = Token.natural    lexer -- parses a natural number
integer :: Parser Integer
integer    = Token.integer    lexer -- parses an integer
comma :: Parser String
comma      = Token.comma       lexer -- parses a comma
whiteSpace :: Parser ()
whiteSpace = Token.whiteSpace lexer -- parses whitespace


optionalCommaSep :: Parser a -> Parser [a]
optionalCommaSep p = many (p <* optional comma)

singleOrList :: Parser a -> Parser [a]
singleOrList p = braces (optionalCommaSep p) <|> (:[]) <$> p

numberOption :: Parser Integer
numberOption = option 1 (brackets natural)

ident :: Parser String
ident = (identifier <|> stringLiteral) <?> "identifier"

identList :: Parser [String]
identList = singleOrList ident

binary :: String -> (a -> a -> a) -> Assoc -> Operator String () Identity a
binary name fun = Infix  ( reservedOp name *> return fun )
prefix :: String -> (a -> a) -> Operator String () Identity a
prefix name fun = Prefix ( reservedOp name *> return fun )

termOperatorTable :: [[Operator String () Identity (Term String)]]
termOperatorTable =
        [ [ prefix "-" Minus ]
        , [ binary "*" (:*:) AssocLeft ]
        , [ binary "+" (:+:) AssocLeft, binary "-" (:-:) AssocLeft ]
        ]

termAtom :: Parser (Term String)
termAtom =  (Var <$> ident)
        <|> (Const <$> integer)
        <|> parens term
        <?> "basic term"

term :: Parser (Term String)
term = buildExpressionParser termOperatorTable termAtom <?> "term"

parseOp :: Parser Op
parseOp = (reservedOp "<" *> return Lt) <|>
          (reservedOp "<=" *> return Le) <|>
          (reservedOp "=" *> return Eq) <|>
          (reservedOp "!=" *> return Ne) <|>
          (reservedOp ">" *> return Gt) <|>
          (reservedOp ">=" *> return Ge)

linIneq :: Parser (Formula String)
linIneq = do
        lhs <- term
        op <- parseOp
        rhs <- term
        return (LinearInequation lhs op rhs)

formOperatorTable :: [[Operator String () Identity (Formula String)]]
formOperatorTable =
        [ [ prefix "!" Neg ]
        , [ binary "&&" (:&:) AssocRight ]
        , [ binary "||" (:|:) AssocRight ]
        ]

formAtom :: Parser (Formula String)
formAtom =  try linIneq
        <|> (reserved "true" *> return FTrue)
        <|> (reserved "false" *> return FFalse)
        <|> parens formula
        <?> "basic formula"

formula :: Parser (Formula String)
formula = buildExpressionParser formOperatorTable formAtom <?> "formula"

quantFormula :: Parser (QuantFormula String)
quantFormula =
        (do
            reserved "EXISTS"
            xs <- optionalCommaSep ident
            reservedOp ":"
            p <- formula
            return (ExQuantFormula xs p)
        )
          <|>
        (ExQuantFormula [] <$> formula)

instance FromJSON (QuantFormula String) where
          parseJSON (String v) = do
                  let formula = parse quantFormula "" (T.unpack v)
                  case formula of
                    Left e -> fail "Predicate formula not well-formed."
                    Right r -> return r
          parseJSON _ = fail "Expect string for predicate."

instance ToJSON (QuantFormula String) where
        toJSON x = String ""

data RecordTransition = RecordTransition { 
        name :: String,
        pre :: [String],
        post :: [String]
} deriving (Show)

data RecordPP = RecordPP { 
        title :: String,
        states :: [String],
        transitions :: [RecordTransition],
        initialStates :: [String],
        trueStates :: [String],
        predicate :: Maybe (QuantFormula String)
} deriving (Show)

$(deriveJSON defaultOptions ''RecordTransition)
$(deriveJSON defaultOptions ''RecordPP)

recordPP2PopulationProtocol :: RecordPP -> PopulationProtocol
recordPP2PopulationProtocol r = 
  makePopulationProtocolFromStrings (title r) (states r) (map name (transitions r)) (initialStates r) (trueStates r) falseStates p arcs where
        falseStates = [q | q <- states r, not (S.member q (S.fromList (trueStates r)))]
        count = \x -> fromIntegral . length . (filter (== x))
        arcs = [(q, name t, (count q) (pre t)) |  t <- transitions r, q <- pre t] ++ 
               [(name t, q, (count q) (post t)) | t <- transitions r, q <- post t]
        p = case predicate r of Nothing -> ExQuantFormula [] FTrue 
                                (Just p') -> p'

parseContent :: Parser PopulationProtocol
parseContent  = do 
  str <- manyTill anyChar eof
  let r = eitherDecode (BS.pack str)
  case r of
    (Left e) ->  fail e
    (Right pp) -> return (recordPP2PopulationProtocol pp)