import Cegar.Core

namespace Cegar.Layered

universe u

inductive Layout (S : Type u) where
  | first (claim : S -> Bool)
  | next (previous : Layout S) (edge : S -> S -> Bool)
      (allowed : S -> Bool) (claim : S -> Bool)

def Layout.claim {S : Type u} : Layout S -> S -> Bool
  | .first p => p
  | .next _ _ _ p => p

inductive Travel {S : Type u} (initial : S -> Bool) : Layout S -> S -> Type u
  | start (p : S -> Bool) (x : S) :
      initial x = true -> Travel initial (.first p) x
  | step {previous : Layout S} {x : S}
      (earlier : Travel initial previous x) (edge : S -> S -> Bool)
      (allowed claim : S -> Bool) (y : S) :
      edge x y = true -> allowed y = true ->
      Travel initial (.next previous edge allowed claim) y

def Conditions {S : Type u} (initial : S -> Bool) : Layout S -> Prop
  | .first p => forall  x, initial x = true -> p x = true
  | .next previous edge allowed p =>
    Conditions initial previous /\
      forall  x y, previous.claim x = true -> edge x y = true ->
        allowed y = true -> p y = true

def checkConditions {S : Type u} (enum : Enumeration S)
    (initial : S -> Bool) : Layout S -> Bool
  | .first p => enum.items.all fun x => !(initial x) || p x
  | .next previous edge allowed p =>
    checkConditions enum initial previous &&
      enum.items.all (fun x => enum.items.all fun y =>
        !(previous.claim x && edge x y && allowed y) || p y)

private theorem impliesBool (p q : Bool) : ((!p || q) = true) <-> (p = true -> q = true) := by
  cases p <;> simp

theorem conditions_reflect {S : Type u} (enum : Enumeration S)
    (initial : S -> Bool) (layout : Layout S) :
    checkConditions enum initial layout = true <-> Conditions initial layout := by
  induction layout with
  | first p => simp only [checkConditions, Conditions, List.all_eq_true, enum.complete, true_implies,
      impliesBool]
  | next previous edge allowed p ih =>
    simp only [checkConditions, Conditions, Bool.and_eq_true, List.all_eq_true,
      enum.complete, true_implies, impliesBool, and_imp, ih]

def check {S : Type u} (enum : Enumeration S) (initial bad : S -> Bool)
    (layout : Layout S) : Bool :=
  checkConditions enum initial layout &&
    enum.items.all (fun x => !(layout.claim x && bad x))

private theorem incompatible (a b : Bool) :
    (Bool.not (a && b) = true) <-> (a = true -> b != true) := by
  cases a <;> cases b <;> decide

theorem check_reflect {S : Type u} (enum : Enumeration S)
    (initial bad : S -> Bool) (layout : Layout S) :
    check enum initial bad layout = true <->
      Conditions initial layout /\ forall  x, layout.claim x = true -> bad x != true := by
  simp only [check, Bool.and_eq_true, conditions_reflect, List.all_eq_true,
    enum.complete, true_implies]
  constructor
  . rintro <h, final>
    exact <h, fun x => (incompatible _ _).mp (final x)>
  . rintro <h, final>
    exact <h, fun x => (incompatible _ _).mpr (final x)>

theorem invariant {S : Type u} {initial : S -> Bool} {layout : Layout S}
    {x : S} (p : Travel initial layout x) (conditions : Conditions initial layout) :
    layout.claim x = true := by
  induction p with
  | start p x hi => exact conditions x hi
  | step earlier edge allowed claim y he ha ih =>
    exact conditions.2 _ y (ih conditions.1) he ha

theorem check_sound {S : Type u} (enum : Enumeration S)
    {initial bad : S -> Bool} {layout : Layout S}
    (verified : check enum initial bad layout = true) :
    (Sigma  x, { _p : Travel initial layout x // bad x = true }) -> False := by
  intro <x, p, hb>
  obtain <hc, hf> := (check_reflect enum initial bad layout).mp verified
  exact hf x (invariant p hc) hb

end Cegar.Layered
