import Cegar.Core

namespace Cegar.Logic

universe u

inductive Term (S : Type u) (n : Nat) where
  | var : Fin n -> Term S n
  | val : S -> Term S n

inductive Formula (S : Type u) (k : Nat) : Nat -> Type u
  | top {n} : Formula S k n
  | bot {n} : Formula S k n
  | atom {n} : Fin k -> Term S n -> Formula S k n
  | equal {n} : Term S n -> Term S n -> Formula S k n
  | edge {n} : Term S n -> Term S n -> Formula S k n
  | both {n} : Formula S k n -> Formula S k n -> Formula S k n
  | either {n} : Formula S k n -> Formula S k n -> Formula S k n
  | impl {n} : Formula S k n -> Formula S k n -> Formula S k n
  | someState {n} : Formula S k (n + 1) -> Formula S k n
  | everyState {n} : Formula S k (n + 1) -> Formula S k n

def Term.eval {S : Type u} {n : Nat} (env : Fin n -> S) : Term S n -> S
  | .var i => env i
  | .val x => x

def extend {S : Type u} {n : Nat} (x : S) (env : Fin n -> S) : Fin (n + 1) -> S :=
  Fin.cases x env

structure Model (S : Type u) (k : Nat) where
  states : Enumeration S
  atom : Fin k -> S -> Bool
  edge : S -> S -> Bool

def Denote {S : Type u} {k : Nat} (M : Model S k) :
    {n : Nat} -> Formula S k n -> (Fin n -> S) -> Prop
  | _, .top, _ => True
  | _, .bot, _ => False
  | _, .atom i t, rho => M.atom i (t.eval rho) = true
  | _, .equal s t, rho => s.eval rho = t.eval rho
  | _, .edge s t, rho => M.edge (s.eval rho) (t.eval rho) = true
  | _, .both p q, rho => Denote M p rho /\ Denote M q rho
  | _, .either p q, rho => Denote M p rho \/ Denote M q rho
  | _, .impl p q, rho => Denote M p rho -> Denote M q rho
  | _, .someState p, rho => exists  x, Denote M p (extend x rho)
  | _, .everyState p, rho => forall  x, Denote M p (extend x rho)

def eval {S : Type u} [DecidableEq S] {k : Nat} (M : Model S k) :
    {n : Nat} -> Formula S k n -> (Fin n -> S) -> Bool
  | _, .top, _ => true
  | _, .bot, _ => false
  | _, .atom i t, rho => M.atom i (t.eval rho)
  | _, .equal s t, rho => decide (s.eval rho = t.eval rho)
  | _, .edge s t, rho => M.edge (s.eval rho) (t.eval rho)
  | _, .both p q, rho => eval M p rho && eval M q rho
  | _, .either p q, rho => eval M p rho || eval M q rho
  | _, .impl p q, rho => !(eval M p rho) || eval M q rho
  | _, .someState p, rho => M.states.items.any fun x => eval M p (extend x rho)
  | _, .everyState p, rho => M.states.items.all fun x => eval M p (extend x rho)

theorem reflects {S : Type u} [DecidableEq S] {k n : Nat} (M : Model S k)
    (p : Formula S k n) (rho : Fin n -> S) : eval M p rho = true <-> Denote M p rho := by
  induction p with
  | top => simp [eval, Denote]
  | bot => simp [eval, Denote]
  | atom => simp [eval, Denote]
  | equal => simp [eval, Denote]
  | edge => simp [eval, Denote]
  | both _ _ hp hq => simp [eval, Denote, hp, hq]
  | either _ _ hp hq => simp [eval, Denote, hp, hq]
  | impl p _ hp hq =>
    simp only [eval, Denote, <- hp, <- hq]
    cases h : eval M p rho <;> simp
  | someState _ hp => simp [eval, Denote, M.states.complete, hp]
  | everyState _ hp => simp [eval, Denote, M.states.complete, hp]

def Term.rename {S : Type u} {n m : Nat} (f : Fin n -> Fin m) :
    Term S n -> Term S m
  | .var i => .var (f i)
  | .val x => .val x

def under {n m : Nat} (f : Fin n -> Fin m) : Fin (n + 1) -> Fin (m + 1) :=
  Fin.cases 0 (fun i => (f i).succ)

def Formula.rename {S : Type u} {k : Nat} :
    {n m : Nat} -> (Fin n -> Fin m) -> Formula S k n -> Formula S k m
  | _, _, _, .top => .top
  | _, _, _, .bot => .bot
  | _, _, f, .atom i t => .atom i (t.rename f)
  | _, _, f, .equal s t => .equal (s.rename f) (t.rename f)
  | _, _, f, .edge s t => .edge (s.rename f) (t.rename f)
  | _, _, f, .both p q => .both (p.rename f) (q.rename f)
  | _, _, f, .either p q => .either (p.rename f) (q.rename f)
  | _, _, f, .impl p q => .impl (p.rename f) (q.rename f)
  | _, _, f, .someState p => .someState (p.rename (under f))
  | _, _, f, .everyState p => .everyState (p.rename (under f))

@[simp] theorem rename_term {S : Type u} {n m : Nat} (f : Fin n -> Fin m)
    (rho : Fin m -> S) (t : Term S n) :
    (t.rename f).eval rho = t.eval (rho  o  f) := by
  cases t <;> rfl

theorem extend_under {S : Type u} {n m : Nat} (f : Fin n -> Fin m)
    (rho : Fin m -> S) (x : S) : extend x rho  o  under f = extend x (rho  o  f) := by
  funext i
  refine Fin.cases ?_ (fun j => ?_) i
  . simp [extend, under, Function.comp_def]
  . simp [extend, under, Function.comp_def]

theorem rename_denote {S : Type u} {k n : Nat} (M : Model S k)
    (p : Formula S k n) {m : Nat} (f : Fin n -> Fin m) (rho : Fin m -> S) :
    Denote M (p.rename f) rho <-> Denote M p (rho  o  f) := by
  induction p generalizing m with
  | top => rfl
  | bot => rfl
  | atom => simp [Formula.rename, Denote]
  | equal => simp [Formula.rename, Denote]
  | edge => simp [Formula.rename, Denote]
  | both _ _ hp hq => simp [Formula.rename, Denote, hp, hq]
  | either _ _ hp hq => simp [Formula.rename, Denote, hp, hq]
  | impl _ _ hp hq => simp [Formula.rename, Denote, hp, hq]
  | someState _ hp => simp [Formula.rename, Denote, hp, extend_under]
  | everyState _ hp => simp [Formula.rename, Denote, hp, extend_under]

theorem rename_eval {S : Type u} [DecidableEq S] {k n m : Nat} (M : Model S k)
    (p : Formula S k n) (f : Fin n -> Fin m) (rho : Fin m -> S) :
    eval M (p.rename f) rho = eval M p (rho  o  f) :=
  Bool.eq_iff_iff.mpr ((reflects _ _ _).trans ((rename_denote M p f rho).trans (reflects _ _ _).symm))

end Cegar.Logic
