module SMT.Core where

open import Agda.Builtin.String
  using (String; primShowNat; primStringAppend)

open import Cubical.Foundations.Prelude
open import Cubical.Data.Bool.Base
  using (Bool; true; false; not; _and_; _or_; if_then_else_; Dec→Bool; _≟_)
open import Cubical.Data.Empty.Base using (⊥)
open import Cubical.Data.Int.Base as ℤ
  using (ℤ; pos; negsuc)
open import Cubical.Data.Int.Properties
  using (discreteℤ)
open import Cubical.Data.Int.Order
  using (≤Dec; <Dec)
open import Cubical.Data.List.Base
  using (List; []; _∷_)
open import Cubical.Data.Nat.Base
  using (suc)
open import Cubical.Data.Sigma.Base
  using (_×_; _,_)
open import Cubical.Data.Unit.Base
  using (Unit; tt)
open import Cubical.Relation.Nullary.Base
  using (¬_)

infixr 5 _++s_
infixr 5 _all∷_

_++s_ : String → String → String
_++s_ = primStringAppend

data Sort : Type₀ where
  sBool : Sort
  sInt  : Sort

Value : Sort → Type₀
Value sBool = Bool
Value sInt = ℤ

record Decl : Type₀ where
  constructor decl
  field
    name : String
    sort : Sort

open Decl public

Ctx : Type₀
Ctx = List Decl

data Var : Ctx → Sort → Type₀ where
  vz : ∀ {Γ x s} → Var (decl x s ∷ Γ) s
  vs : ∀ {Γ x s t} → Var Γ s → Var (decl x t ∷ Γ) s

Env : Ctx → Type₀
Env [] = Unit
Env (d ∷ Γ) = Value (sort d) × Env Γ

lookupEnv : ∀ {Γ s} → Var Γ s → Env Γ → Value s
lookupEnv vz (v , _) = v
lookupEnv (vs x) (_ , ρ) = lookupEnv x ρ

varSMT : ∀ {Γ s} → Var Γ s → String
varSMT {Γ = decl x _ ∷ _} vz = x
varSMT (vs x) = varSMT x

data Expr (Γ : Ctx) : Sort → Type₀ where
  bool   : Bool → Expr Γ sBool
  int    : ℤ → Expr Γ sInt
  var    : ∀ {s} → Var Γ s → Expr Γ s

  iNeg   : Expr Γ sInt → Expr Γ sInt
  iAdd   : Expr Γ sInt → Expr Γ sInt → Expr Γ sInt
  iSub   : Expr Γ sInt → Expr Γ sInt → Expr Γ sInt
  iScale : ℤ → Expr Γ sInt → Expr Γ sInt

  iEq    : Expr Γ sInt → Expr Γ sInt → Expr Γ sBool
  iNe    : Expr Γ sInt → Expr Γ sInt → Expr Γ sBool
  iLe    : Expr Γ sInt → Expr Γ sInt → Expr Γ sBool
  iLt    : Expr Γ sInt → Expr Γ sInt → Expr Γ sBool
  iGe    : Expr Γ sInt → Expr Γ sInt → Expr Γ sBool
  iGt    : Expr Γ sInt → Expr Γ sInt → Expr Γ sBool

  bNot   : Expr Γ sBool → Expr Γ sBool
  bAnd   : Expr Γ sBool → Expr Γ sBool → Expr Γ sBool
  bOr    : Expr Γ sBool → Expr Γ sBool → Expr Γ sBool
  bXor   : Expr Γ sBool → Expr Γ sBool → Expr Γ sBool
  bImp   : Expr Γ sBool → Expr Γ sBool → Expr Γ sBool
  bEq    : Expr Γ sBool → Expr Γ sBool → Expr Γ sBool

  ite    : ∀ {s} → Expr Γ sBool → Expr Γ s → Expr Γ s → Expr Γ s

eval : ∀ {Γ s} → Env Γ → Expr Γ s → Value s
eval ρ (bool b) = b
eval ρ (int z) = z
eval ρ (var x) = lookupEnv x ρ
eval ρ (iNeg x) = ℤ.-_ (eval ρ x)
eval ρ (iAdd x y) = ℤ._+_ (eval ρ x) (eval ρ y)
eval ρ (iSub x y) = ℤ._-_ (eval ρ x) (eval ρ y)
eval ρ (iScale k x) = ℤ._·_ k (eval ρ x)
eval ρ (iEq x y) = Dec→Bool (discreteℤ (eval ρ x) (eval ρ y))
eval ρ (iNe x y) = not (Dec→Bool (discreteℤ (eval ρ x) (eval ρ y)))
eval ρ (iLe x y) = Dec→Bool (≤Dec (eval ρ x) (eval ρ y))
eval ρ (iLt x y) = Dec→Bool (<Dec (eval ρ x) (eval ρ y))
eval ρ (iGe x y) = Dec→Bool (≤Dec (eval ρ y) (eval ρ x))
eval ρ (iGt x y) = Dec→Bool (<Dec (eval ρ y) (eval ρ x))
eval ρ (bNot x) = not (eval ρ x)
eval ρ (bAnd x y) = eval ρ x and eval ρ y
eval ρ (bOr x y) = eval ρ x or eval ρ y
eval ρ (bXor x y) = if eval ρ x then not (eval ρ y) else eval ρ y
eval ρ (bImp x y) = not (eval ρ x) or eval ρ y
eval ρ (bEq x y) = Dec→Bool ((eval ρ x) ≟ (eval ρ y))
eval ρ (ite c t f) = if eval ρ c then eval ρ t else eval ρ f

sortSMT : Sort → String
sortSMT sBool = "Bool"
sortSMT sInt = "Int"

showInt : ℤ → String
showInt (pos n) = primShowNat n
showInt (negsuc n) = "(- " ++s primShowNat (suc n) ++s ")"

paren1 : String → String → String
paren1 op x = "(" ++s op ++s " " ++s x ++s ")"

paren2 : String → String → String → String
paren2 op x y = "(" ++s op ++s " " ++s x ++s " " ++s y ++s ")"

paren3 : String → String → String → String → String
paren3 op x y z =
  "(" ++s op ++s " " ++s x ++s " " ++s y ++s " " ++s z ++s ")"

exprSMT : ∀ {Γ s} → Expr Γ s → String
exprSMT (bool true) = "true"
exprSMT (bool false) = "false"
exprSMT (int z) = showInt z
exprSMT (var x) = varSMT x
exprSMT (iNeg x) = paren1 "-" (exprSMT x)
exprSMT (iAdd x y) = paren2 "+" (exprSMT x) (exprSMT y)
exprSMT (iSub x y) = paren2 "-" (exprSMT x) (exprSMT y)
exprSMT (iScale k x) = paren2 "*" (showInt k) (exprSMT x)
exprSMT (iEq x y) = paren2 "=" (exprSMT x) (exprSMT y)
exprSMT (iNe x y) = paren2 "distinct" (exprSMT x) (exprSMT y)
exprSMT (iLe x y) = paren2 "<=" (exprSMT x) (exprSMT y)
exprSMT (iLt x y) = paren2 "<" (exprSMT x) (exprSMT y)
exprSMT (iGe x y) = paren2 ">=" (exprSMT x) (exprSMT y)
exprSMT (iGt x y) = paren2 ">" (exprSMT x) (exprSMT y)
exprSMT (bNot x) = paren1 "not" (exprSMT x)
exprSMT (bAnd x y) = paren2 "and" (exprSMT x) (exprSMT y)
exprSMT (bOr x y) = paren2 "or" (exprSMT x) (exprSMT y)
exprSMT (bXor x y) = paren2 "xor" (exprSMT x) (exprSMT y)
exprSMT (bImp x y) = paren2 "=>" (exprSMT x) (exprSMT y)
exprSMT (bEq x y) = paren2 "=" (exprSMT x) (exprSMT y)
exprSMT (ite c t f) = paren3 "ite" (exprSMT c) (exprSMT t) (exprSMT f)

declSMT : Decl → String
declSMT d =
  "(declare-const " ++s name d ++s " " ++s sortSMT (sort d) ++s ")"

declsSMT : Ctx → String
declsSMT [] = ""
declsSMT (d ∷ Γ) = declSMT d ++s "\n" ++s declsSMT Γ

assertionsSMT : ∀ {Γ} → List (Expr Γ sBool) → String
assertionsSMT [] = ""
assertionsSMT (p ∷ ps) =
  "(assert " ++s exprSMT p ++s ")\n" ++s assertionsSMT ps

record Problem (Γ : Ctx) : Type₀ where
  constructor problem
  field
    assumptions : List (Expr Γ sBool)
    claim       : Expr Γ sBool

open Problem public

counterexampleCheckSMT : ∀ {Γ} → Problem Γ → String
counterexampleCheckSMT {Γ = Γ} p =
  "(set-logic QF_LIA)\n" ++s
  declsSMT Γ ++s
  assertionsSMT (assumptions p) ++s
  "(assert " ++s exprSMT (bNot (claim p)) ++s ")\n" ++s
  "(check-sat)\n"

counterexampleQuerySMT : ∀ {Γ} → Problem Γ → String
counterexampleQuerySMT p =
  counterexampleCheckSMT p ++s
  "(get-model)\n"

data AllTrue {Γ : Ctx} : List (Expr Γ sBool) → Env Γ → Type₀ where
  all[]   : ∀ {ρ} → AllTrue [] ρ
  _all∷_ : ∀ {ρ p ps} → eval ρ p ≡ true → AllTrue ps ρ → AllTrue (p ∷ ps) ρ

Statement : ∀ {Γ} → Problem Γ → Type₀
Statement {Γ = Γ} p =
  (ρ : Env Γ) → AllTrue (assumptions p) ρ → eval ρ (claim p) ≡ true

record Counterexample {Γ : Ctx} (p : Problem Γ) : Type₀ where
  constructor counterexample
  field
    env             : Env Γ
    assumptionsTrue : AllTrue (assumptions p) env
    claimFalse      : eval env (claim p) ≡ false

open Counterexample public

false≠true : false ≡ true → ⊥
false≠true p = subst (λ b → if b then ⊥ else Unit) p tt

refute : ∀ {Γ} {p : Problem Γ} → Counterexample p → ¬ Statement p
refute c valid =
  false≠true (sym (claimFalse c) ∙ valid (env c) (assumptionsTrue c))