module Verify (verify) where

import Control.Monad.State hiding (State)
import Data.List
import Math.SMT.Yices.Pipe
import Math.SMT.Yices.Syntax
import System.IO
import Text.Printf

import Error
import Model

-- | Verify a model with k-induction.
verify :: FilePath -> Int -> Model -> IO ()
verify _ maxK _ | maxK < 1 = error "max k can not be less than 1"
verify yices maxK m = do
  mapM_ (verifyModel yices format maxK) models
  where
  model = typeCheckModel m
  assertions = [ a | a@(Assert _ _) <- actions model ]
  models = map (trimModel . trimAssertions model) assertions
  format = "verifying %-" ++ show (maximum [ length (intercalate "." n) | (Assert _ n) <- actions model ]) ++ "s    "

-- | Remove all assertions except one.
trimAssertions :: Model -> Action -> Model
trimAssertions m a = m { actions = filter keep $ actions m }
  where
  keep b@(Assert _ _) = a == b
  keep _ = True

-- | Trim all unneeded stuff from a model.
trimModel :: Model -> Model
trimModel = id  --XXX

-- | Verify a trimmed model.
verifyModel :: FilePath -> String -> Int -> Model -> IO ()
verifyModel y format maxK m = do
  printf format (intercalate "." path) >> hFlush stdout
  env0 <- initEnv y m
  execStateT (check env0 1) env0
  return ()
  where
  [path] = [ path | Assert _ path <- actions m ]
  check :: Env -> Int -> Y ()
  check _ k | k > maxK = liftIO $ printf "inconclusive: unable to proved step up to max k = %d\n" maxK
  check env0 k = do
    transition m
    resultBasis <- checkBasis m env0
    case resultBasis of
      Fail    -> liftIO $ printf "FAILED: disproved basis in k = %d\n" k
      Problem -> return ()
      Pass -> do
        resultStep <- checkStep
        case resultStep of
          Fail    -> check env0 (k + 1)
          Problem -> return ()
          Pass    -> liftIO $ printf "passed: proved step in k = %d\n" k
        
data Result = Pass | Fail | Problem

-- | Check induction step.
checkStep :: Y Result
checkStep = do
  Env y _ _ cmds asserts <- get
  r <- liftIO $ quickCheckY y [] $ reverse $ [ASSERT $ NOT $ head asserts] ++ cmds
  return $ result r

-- | Check induction basis.
checkBasis :: Model -> Env -> Y Result
checkBasis m env0 = do
  Env y _ _ cmds asserts <- get
  r <- liftIO $ quickCheckY y [] $ reverse $ [ASSERT $ NOT $ AND asserts] ++ [ ASSERT $ VarE (getVar' env0 (State a)) := exprConst t c | a@(VS _ t c _) <- variables m ] ++ cmds
  return $ result r

result :: ResY -> Result
result a = case a of
  Sat _   -> Fail
  UnSat _ -> Pass
  InCon _ -> Problem
  _ -> error $ "unexpected yices results: " ++ show a


-- | Transition relation.  Returns assertion.
transition :: Model -> Y ()
transition m = mapM_ action $ actions m

action :: Action -> Y ()
action a = case a of
  Assign v e -> do
    e <- expr (typeOf v) e
    v <- addVar v
    Env y i table cmds asserts <- get
    put $ Env y i table (ASSERT (VarE v := e) : cmds) asserts
  Assert a _ -> do
    a <- expr Bool a
    Env y i table cmds asserts <- get
    put $ Env y i table cmds (a : asserts)
  Assume a _ -> do
    a <- expr Bool a
    Env y i table cmds asserts <- get
    put $ Env y i table (ASSERT a : cmds) asserts

expr :: Type -> E -> Y ExpY
expr t a = case a of
  Var a@(Volatile _ _ _) -> addVar a >>= return . VarE
  Var a   -> getVar a >>= return . VarE
  Const a -> return $ exprConst t a
  Not a _ -> expr Bool a >>= return . NOT

  And a b _ -> do
    a <- expr Bool a
    b <- expr Bool b
    return $ AND [a, b]

  Or  a b _ -> do
    a <- expr Bool a
    b <- expr Bool b
    return $ OR  [a, b]

  Mul a b _ -> do
    a <- expr t a
    b <- expr t b
    return $ a :*: b
  
  Div a b _  -> do
    a <- expr t a
    b <- expr t b
    return $ op a b
    where
    op = case t of
      Integer _ -> DIV
      _         -> (:/:)

  Mod a b _ -> do
    a <- expr t a
    b <- expr t b
    return $ MOD a b

  Add a b _ -> do
    a <- expr t a
    b <- expr t b
    return $ a :+: b

  Sub a b _ -> do
    a <- expr t a
    b <- expr t b
    return $ a :-: b

  Lt  a b n -> do
    let t' = unify n (typeOf a) (typeOf b)
    a <- expr t' a
    b <- expr t' b
    return $ a :< b

  Eq  a b n -> do
    let t' = unify n (typeOf a) (typeOf b)
    a <- expr t' a
    b <- expr t' b
    return $ a := b

  Mux a b c n -> do
    let t' = unify n (typeOf b) (typeOf c)
    a <- expr Bool a
    b <- expr t' b
    c <- expr t' c
    return $ IF a b c

  Ref _ _ -> notSupported a "address-of operator"
  Deref _ _ -> notSupported a "indirection operator"

exprConst :: Type -> Const -> ExpY
exprConst t a = case a of 
  CBool     a _ -> LitB a
  CInteger  a _ | t == Bool -> LitB $ a /= 0
                | otherwise -> LitI a
  CRational a _ -> LitR a



type Y = StateT Env IO

data Env = Env FilePath Int [(V, String)] [CmdY] [ExpY]  -- Env nextId table cmds assertions

initEnv :: FilePath -> Model -> IO Env
initEnv y m = execStateT (mapM_ addVar $ map State $ variables m) (Env y 0 [] [] [])

addVar :: V -> Y String
addVar v = do
  Env y i table cmds asserts <- get
  let name = printf "n%d" i
  put $ Env y (i + 1) (replace (v, name) table) (DEFINE (name, VarT $ typeY v) Nothing : cmds) asserts
  return name

getVar :: V -> Y String
getVar v = do
  env <- get
  return $ getVar' env v

getVar' :: Env -> V -> String
getVar' (Env _ _ table _ _) v = case lookup v table of
  Just a -> a
  Nothing -> error $ "Verify.getVar: variable not found: " ++ show v  ++ " in " ++ show (fst $ unzip table)

replace :: Eq a => (a, b) -> [(a, b)] -> [(a, b)]
replace a [] = [a]
replace (a, b) ((a', b') : c) | a == a'   = (a, b) : c
                              | otherwise = (a', b') : replace (a, b) c


typeY :: TypeOf a => a -> String
typeY a = case typeOf a of
  Void       -> error "Verify.typeY: void"
  Ptr _      -> error "Verify.typeY: ptr"
  Bool       -> "bool"
  Integer  _ -> "int"
  Rational _ -> "real"