module SMT where

import Data.List
import Data.Char
import Data.Either
import Text.ParserCombinators.Parsec
import System.IO
import System.Process

data Exp =
    Var String
  | Val Int
  | Plus [Exp]
  | Times [Exp]
  | Ite Formula Exp Exp
  deriving Eq

data Formula =
    And [Formula]
  | Or [Formula]
  | Not Formula
  | Distinct [Exp]
  | Geq Exp Exp
  | Gt Exp Exp
  | Eq Exp Exp
  | Iff Formula Formula
  | FVar String
  deriving Eq

data Command =
    DeclareInt String
  | DeclareNat String
  | DeclareBool String
  | Assert Formula
  | AssertSoft Formula
  | CheckSat
  | GetValue [String]

type SMTInput = [Command]

type Model = [(String, Int)]

type SMTOutput = Maybe Model

-- Escape/unescape functions for identifiers in SMT solvers.

escape' :: String -> String
escape' []         = []
escape' ('_' : cs) = "__" ++ escape' cs
escape' (c : cs)
  | isAlphaNum c = c : escape' cs
  | otherwise    = "_x" ++ show (fromEnum c) ++ "_" ++ escape' cs

escape :: String -> String
escape s = '_': escape' s

unescape' :: String -> String
unescape' [] = []
unescape' ('_' : 'x' : cs)
  | (s, _ : cs') <- break (== '_') cs = chr (read s :: Int) : unescape' cs'
unescape' "_"             = error "unexpacted end of string"
unescape' ('_': '_' : cs) = '_' : unescape' cs
unescape' ('_': c : _)    = error ("unknown escape code: " ++ [c])
unescape' (c : cs)        = c : unescape' cs

unescape :: String -> String
unescape ('_' : cs) = unescape' cs
unescape s          = error ("unexpected string: " ++ show s)


{-
escape :: String -> String
escape s = "|" ++ s ++ "|"

unescape :: String -> String
unescape s = [ c | (i, c) <- zip [0..] s, 0 < i && i < n-2 ]
  where n = length s
-}

-- Pretty printers.

showApplication :: Show a => String -> [a] -> String
showApplication s xs =
  "(" ++ intercalate " " (s : [ show x | x <- xs ]) ++ ")"

instance Show Exp where
  show (Val n)       = show n
  show (Var x)       = escape x
  show (Plus [])     = "0"
  show (Plus [e])    = show e
  show (Plus es)     = showApplication "+" es
  show (Times [])    = "1"
  show (Times [e])   = show e
  show (Times es)    = showApplication "*" es
  show (Ite f e1 e2) = "(ite " ++ show f ++ " " ++ show e1 ++ " " ++ show e2 ++ ")"

instance Show Formula where
  show (And [])         = "true"
  show (Or [])          = "false"
  show (And [f])        = show f
  show (Or [f])         = show f
  show (And fs)         = showApplication "and" fs
  show (Or fs)          = showApplication "or"  fs
  show (Not f)          = showApplication "not" [f]
  show (Distinct [])    = "true"
  show (Distinct [_])   = "true"
  show (Distinct es)    = showApplication "distinct" es
  show (Gt e1 e2)       = showApplication ">"   [e1,e2]
  show (Geq e1 e2)      = showApplication ">="  [e1,e2]
  show (Eq e1 e2)       = showApplication "="   [e1,e2]
  show (Iff f1 f2)      = showApplication "="   [f1,f2]
  show (FVar x)         = escape x
--  show (BoolVar x)      = escape x

instance Show Command where
  show (DeclareInt x)  = "(declare-fun " ++ escape x ++ " () Int)"
  show (DeclareNat x)  = show (DeclareInt x) ++ show (Assert (Var x`Geq` Val 0))
  show (DeclareBool x) = "(declare-fun " ++ escape x ++ " () Bool)"
  show (Assert f)      = showApplication "assert" [f]
  show (AssertSoft f)  = showApplication "assert-soft" [f]
  show (CheckSat)      = "(check-sat)"
  show (GetValue xs)   = "(get-value (" ++ intercalate " " [ escape x | x <- xs ] ++ "))"

showSMTInput :: [Command] -> String
showSMTInput cs = unlines [ show c | c <- cs ]


-- Parsing outputs of SMT solvers.

keyword :: String -> Parser ()
keyword s = do
  spaces
  _ <- string s
  spaces

parseVar :: Parser String
parseVar = do
  spaces
  s <- many1 (noneOf " \n\r\t\0,()")
  spaces
  return (unescape s)

parseNegative :: Parser Int
parseNegative = do
  keyword "("
  keyword "-"
  t <- many1 digit
  keyword ")"
  return (- (read t :: Int))

parseNonnegative :: Parser Int
parseNonnegative = do
  spaces
  t <- many1 digit
  spaces
  return (read t :: Int)

parseInt :: Parser Int
parseInt = do
  i <- parseNegative <|> parseNonnegative
  return i

parseTrue :: Parser Int
parseTrue = do
  keyword "true"
  return 1

parseFalse :: Parser Int
parseFalse = do
  keyword "false"
  return 0

parseBool :: Parser Int
parseBool = parseTrue <|> parseFalse

parseValue :: Parser Int
parseValue = try parseBool <|> try parseInt

parsePair :: Parser (String, Int)
parsePair = do
  keyword "("
  x <- parseVar
  y <- parseValue
  keyword ")"
  return (x, y)

parseModel :: Parser Model
parseModel = do
  keyword "("
  m <- many parsePair
  keyword ")"
  return m

parseSat :: Parser SMTOutput
parseSat = do
  keyword "sat"
  m <- option [] parseModel
  return (Just m)

parseUnsat :: Parser SMTOutput
parseUnsat = do
  keyword "unsat"
  return Nothing

parseSMTOutput :: Parser SMTOutput
parseSMTOutput = try parseSat <|> parseUnsat

-- Constructors.

plus' :: [Exp] -> [Exp]
plus' []              = []
plus' (Val 0 : es)    = plus' es
plus' (Plus es' : es) = plus' (es' ++ es)
plus' (e : es)        = e : plus' es

plus :: [Exp] -> Exp
plus es =
  case plus' es of
    []  -> Val 0
    es' -> Plus es'

times' :: [Exp] -> [Exp]
times' []              = []
times' (Times es' : es) = times' (es' ++ es)
times' (e : es)        = e : times' es

value :: Exp -> Bool
value (Val _) = True
value _       = False

times :: [Exp] -> Exp
times es
  | k == 0    = Val 0
  | k == 1    = if es2 == [] then Val 1 else Times es2
  | otherwise = Times (Val k : es2)
  where
    es1 = times' es
    k   = product [ n | Val n <- es1 ]
    es2 = [ e | e <- es1, not (value e) ]

top :: Formula
top = And []

bottom :: Formula
bottom = Or []

formula :: Bool -> Formula
formula True  = top
formula False = bottom

neg :: Formula -> Formula
neg (Or [])  = And []
neg (And []) = Or []
neg (Not f)  = f
neg f        = Not f

conj' :: [Formula] -> [Formula] -> Formula
conj' [f] []              = f
conj' fs []               = And fs
conj' fs (And fs' : fs'') = conj' fs (fs' ++ fs'')
conj' _  (Or [] : _)      = bottom
conj' fs (f : fs')        = conj' (f : fs) fs'

conj :: [Formula] -> Formula
conj fs = conj' [] fs

disj' :: [Formula] -> [Formula] -> Formula
disj' [f] []             = f
disj' fs []              = Or fs
disj' fs (Or fs' : fs'') = disj' fs (fs' ++ fs'')
disj' _  (And [] : _)    = top
disj' fs (f : fs')       = disj' (f : fs) fs'

disj :: [Formula] -> Formula
disj fs = disj' [] fs

implies :: Formula -> Formula -> Formula
implies f g = disj [neg f, g]

iff :: Formula -> Formula -> Formula
iff f (And []) = f
iff (And []) f = f
iff f (Or [])  = neg f
iff (Or []) f  = neg f
iff f1      f2 = Iff f1 f2

distinct :: [Exp] -> Formula
distinct []  = top
distinct [_] = top
distinct es  = Distinct es

ite :: Formula -> Exp -> Exp -> Exp
ite (And []) e _ = e
ite (Or  []) _ e = e
ite f1 (Ite f2 e11 e12) e2
  | e12 == e2    = ite (conj [f1, f2]) e11 e2
ite f e1 e2      = Ite f e1 e2

-- Queries

sizeExp :: Exp -> Int
sizeExp (Val _)       = 1
sizeExp (Var _)       = 1
sizeExp (Plus es)     = 1 + sum [ sizeExp e | e <- es ]
sizeExp (Times es)    = 1 + sum [ sizeExp e | e <- es ]
sizeExp (Ite f e1 e2) = 1 + sizeFormula f + sizeExp e1 + sizeExp e2

sizeFormula :: Formula -> Int
sizeFormula (FVar _)      = 1
sizeFormula (Not f)       = 1 + sizeFormula f
sizeFormula (And fs)      = 1 + sum [ sizeFormula f | f <- fs ]
sizeFormula (Or  fs)      = 1 + sum [ sizeFormula f | f <- fs ]
sizeFormula (Iff e1 e2)   = 1 + sizeFormula e1 + sizeFormula e2
sizeFormula (Distinct es) = 1 + sum [ sizeExp e | e <- es ]
sizeFormula (Eq  e1 e2)   = 1 + sizeExp e1 + sizeExp e2
sizeFormula (Geq e1 e2)   = 1 + sizeExp e1 + sizeExp e2
sizeFormula (Gt  e1 e2)   = 1 + sizeExp e1 + sizeExp e2

-- Evaluation functions.

are_distinct :: Eq a => [a] -> Bool
are_distinct []       = True
are_distinct (x : xs) = all (/= x) xs && are_distinct xs

eval_exp :: Model -> Exp -> Int
eval_exp _ (Val n)  = n
eval_exp a (Var x)
  | Just n <- lookup x a = n
  | otherwise            = 0
eval_exp a (Plus es) = sum [ eval_exp a e | e <- es ]
eval_exp a (Times es) = product [ eval_exp a e | e <- es ]
eval_exp a (Ite f e1 e2)
  | eval_formula a f = eval_exp a e1
  | otherwise       = eval_exp a e2

eval_formula :: Model -> Formula -> Bool
eval_formula a (FVar x)      = lookup x a /= Just 0
eval_formula a (And fs)      = and [ eval_formula a f | f <- fs ]
eval_formula a (Or fs)       = or [ eval_formula a f | f <- fs ]
eval_formula a (Not f)       = not (eval_formula a f)
eval_formula a (Iff f1 f2)   = eval_formula a f1 == eval_formula a f2
eval_formula a (Distinct es) = are_distinct [ eval_exp a e | e <- es ]
eval_formula a (Eq e1 e2)    = eval_exp a e1 == eval_exp a e2
eval_formula a (Geq e1 e2)   = eval_exp a e1 >= eval_exp a e2
eval_formula a (Gt e1 e2)    = eval_exp a e1 >  eval_exp a e2

-- Functions for expressions and formulas.

variablesInExp :: Exp -> [Either Exp String]
variablesInExp (Var x)       = [Left (Var x)]
variablesInExp (Val _)       = []
variablesInExp (Plus es)     = variablesInExps es
variablesInExp (Times es)    = variablesInExps es
variablesInExp (Ite f e1 e2) =
  nub (variablesInFormula f ++ variablesInExp e1 ++ variablesInExp e2)

variablesInExps :: [Exp] -> [Either Exp String]
variablesInExps es = nub [ x | e <- es, x <- variablesInExp e ]

variablesInFormula :: Formula -> [Either Exp String]
variablesInFormula (FVar x)      = [Right x]
variablesInFormula (Distinct es) = variablesInExps es
variablesInFormula (Geq e1 e2)   = variablesInExps [e1, e2]
variablesInFormula (Eq e1 e2)    = variablesInExps [e1, e2]
variablesInFormula (Gt e1 e2)    = variablesInExps [e1, e2]
variablesInFormula (Not f)       = variablesInFormula f
variablesInFormula (And fs)      = variablesInFormulas fs
variablesInFormula (Or fs)       = variablesInFormulas fs
variablesInFormula (Iff f1 f2)   = variablesInFormulas [f1, f2]

variablesInFormulas :: [Formula] -> [Either Exp String]
variablesInFormulas fs = nub [ x | f <- fs, x <- variablesInFormula f ]

-- ([Nat Var], [Int Var], [Bool Var])
variables :: Formula -> ([String], [String])
variables f =
  let (vs, bs) = partitionEithers (variablesInFormula f) in
  ([x | Var x <- vs], bs)

variables_in_formulas :: [Formula] -> ([String], [String])
variables_in_formulas [] = ([], [])
variables_in_formulas (f : fs) =
  (union ints1 ints2, union bools1 bools2)
  where
    (ints1, bools1) = variables f
    (ints2, bools2) = variables_in_formulas fs

smt2 :: [Formula] -> Formula -> SMTInput
smt2 softs hard = 
  ints ++ bools ++ 
  [ AssertSoft soft | soft <- softs ] ++ 
  [Assert hard, CheckSat] ++
  if xs == [] && bs == [] then [] else [GetValue (xs ++ bs)]
  where 
    (xs, bs) = variables_in_formulas (hard : softs)
    ints  = [ DeclareInt x | x <- xs ]
    bools = [ DeclareBool b | b <- bs ]

-- Running an SMT solver.

execute :: String -> String -> IO String
execute tool input = do
  (Just hin, Just hout, _, _) <-
    createProcess (proc tool ["/dev/stdin"]) {
      std_in  = CreatePipe,
      std_out = CreatePipe }
  hPutStr hin input
  hClose hin
  hGetContents hout

run :: String -> SMTInput -> IO SMTOutput
run tool input = do
  -- putStrLn ("-- SMT input --\n" ++ showSMTInput input ++ "\n")
  s <- execute tool (showSMTInput input)
  -- putStrLn ("-- SMT output --\n" ++ s ++ "\n")
  case parse parseSMTOutput "(stdin)" s of
    Left e  -> do
      putStrLn (showSMTInput input)
      putStrLn s
      error (show e)
    Right m -> return m

sat :: String -> Formula -> IO SMTOutput
sat tool f = run tool (smt2 [] f)

maxsat :: String -> [Formula] -> Formula -> IO SMTOutput
maxsat tool softs hard = run tool (smt2 softs hard)
