import Std

namespace Cegar

universe u v w

structure Graph (A : Type u) where
  initial : A -> Bool
  bad : A -> Bool
  edge : A -> A -> Bool

inductive Walk {A : Type u} (G : Graph A) : A -> Type u
  | start (a : A) (initial : G.initial a = true) : Walk G a
  | step {a : A} (previous : Walk G a) (b : A) (edge : G.edge a b = true) : Walk G b

def Walk.nodes {A : Type u} {G : Graph A} {a : A} : Walk G a -> List A
  | .start a _ => [a]
  | .step p b _ => p.nodes ++ [b]

theorem Walk.last_mem {A : Type u} {G : Graph A} {a : A} (p : Walk G a) :
    a  in  p.nodes := by
  cases p <;> simp [Walk.nodes]

def BadWalk {A : Type u} (G : Graph A) :=
  Sigma  a, { _p : Walk G a // G.bad a = true }

def Safe {A : Type u} (G : Graph A) : Prop := BadWalk G -> False

inductive Verdict {A : Type u} (G : Graph A) : Type u
  | safe : Safe G -> Verdict G
  | counterexample : BadWalk G -> Verdict G

structure Enumeration (A : Type u) where
  items : List A
  complete : forall  a, a  in  items

def finEnumeration (n : Nat) : Enumeration (Fin n) where
  items := List.finRange n
  complete := fun a => List.mem_ofFn.mpr <a, rfl>

def sumEnumeration {A : Type u} {B : Type v}
    (left : Enumeration A) (right : Enumeration B) : Enumeration (Sum A B) where
  items := left.items.map Sum.inl ++ right.items.map Sum.inr
  complete := by
    intro a
    cases a with
    | inl a => exact List.mem_append_left _ (List.mem_map.mpr <a, left.complete a, rfl>)
    | inr b => exact List.mem_append_right _ (List.mem_map.mpr <b, right.complete b, rfl>)

inductive Pick {A : Type u} {B : A -> Type v} (f : (a : A) -> Option (B a)) :
    List A -> Type (max u v)
  | missing {xs} : (forall  a  in  xs, f a = none) -> Pick f xs
  | found {xs} (a : A) (value : B a) (rest : List A) :
      f a = some value ->
      (forall  b  in  xs, b = a \/ b  in  rest) ->
      rest.length < xs.length ->
      Pick f xs

def pick {A : Type u} {B : A -> Type v} (f : (a : A) -> Option (B a)) :
    (xs : List A) -> Pick f xs
  | [] => .missing (by simp)
  | x :: xs => match h : f x with
    | some value => .found x value xs h (by simp) (by simp)
    | none => match pick f xs with
      | .missing miss => .missing (by
          intro a ha
          rcases List.mem_cons.mp ha with rfl | hm
          . exact h
          . exact miss a hm)
      | .found a value rest hv cover shorter =>
          .found a value (x :: rest) hv (by
            intro b hb
            rcases List.mem_cons.mp hb with rfl | hm
            . exact Or.inr (by simp)
            . rcases cover b hm with he | hr
              . exact Or.inl he
              . exact Or.inr (List.mem_cons_of_mem _ hr)) (by simpa using shorter)

inductive Checked (A : Type u) : Type u
  | yes : A -> Checked A
  | no : (A -> False) -> Checked A

def choose {A : Type u} {B : A -> Type v}
    (enum : Enumeration A) (f : (a : A) -> Checked (B a)) : Checked (Sigma  a, B a) :=
  let option (a : A) : Option (B a) := match f a with
    | .yes b => some b
    | .no _ => none
  match pick option enum.items with
  | .found a b _ _ _ _ => .yes <a, b>
  | .missing miss => .no (fun <a, b> => by
      have h := miss a (enum.complete a)
      cases ha : f a with
      | yes proof => simp [option, ha] at h
      | no refute => exact refute b)

def Prefix {A : Type u} (G : Graph A) := Sigma  a, Walk G a

def extendFrom {A : Type u} (G : Graph A) (b : A) (p : Prefix G) :
    Option (Walk G b) :=
  if h : G.edge p.1 b = true then some (.step p.2 b h) else none

def discover {A : Type u} (G : Graph A) (seen : List (Prefix G)) (b : A) :
    Option (Walk G b) :=
  if h : G.initial b = true then some (.start b h)
  else match pick (extendFrom G b) seen with
    | .found _ p _ _ _ _ => some p
    | .missing _ => none

theorem discover_initial {A : Type u} {G : Graph A} {seen : List (Prefix G)}
    {b : A} (missing : discover G seen b = none) : G.initial b != true := by
  intro h
  simp [discover, h] at missing

theorem discover_closed {A : Type u} {G : Graph A} {seen : List (Prefix G)}
    {b : A} (missing : discover G seen b = none) (p : Prefix G) (hp : p  in  seen) :
    G.edge p.1 b != true := by
  intro h
  unfold discover at missing
  split at missing
  . contradiction
  . split at missing
    . contradiction
    . rename_i hm _
      have impossible := hm p hp
      simp [extendFrom, h] at impossible

def Seen {A : Type u} {G : Graph A} (seen : List (Prefix G)) (a : A) : Prop :=
  exists  p  in  seen, p.1 = a

theorem closedWalk {A : Type u} {G : Graph A} {seen : List (Prefix G)}
    (roots : forall  a, G.initial a = true -> Seen seen a)
    (closed : forall  a b, Seen seen a -> G.edge a b = true -> Seen seen b)
    {a : A} (p : Walk G a) : Seen seen a := by
  induction p with
  | start a hi => exact roots a hi
  | step p b he ih => exact closed _ b ih he

def explore {A : Type u} (G : Graph A) (seen : List (Prefix G))
    (remaining : List A)
    (cover : forall  a, a  in  remaining \/ Seen seen a)
    (good : forall  p  in  seen, G.bad p.1 = false) : Verdict G :=
  match pick (discover G seen) remaining with
  | .missing miss =>
    .safe (fun <a, p, bad> => by
      have roots : forall  b, G.initial b = true -> Seen seen b := by
        intro b hi
        rcases cover b with hm | hs
        . exact (discover_initial (miss b hm) hi).elim
        . exact hs
      have closed : forall  a b, Seen seen a -> G.edge a b = true -> Seen seen b := by
        intro a b hs he
        rcases cover b with hm | found
        . obtain <q, hq, rfl> := hs
          exact (discover_closed (miss b hm) q hq he).elim
        . exact found
      obtain <q, hq, h> := closedWalk roots closed p
      have hb := good q hq
      rw [h, bad] at hb
      contradiction)
  | .found a p rest _ kept shorter =>
    if bad : G.bad a = true then .counterexample <a, p, bad>
    else explore G (<a, p> :: seen) rest (by
      intro b
      rcases cover b with hm | <q, hq, he>
      . rcases kept b hm with rfl | hr
        . exact Or.inr <<b, p>, by simp, rfl>
        . exact Or.inl hr
      . exact Or.inr <q, List.mem_cons_of_mem _ hq, he>) (by
      intro q hq
      rcases List.mem_cons.mp hq with rfl | hm
      . cases h : G.bad a with
        | false => rfl
        | true => exact (bad h).elim
      . exact good q hm)
termination_by remaining.length

def solve {A : Type u} (G : Graph A) (enum : Enumeration A) : Verdict G :=
  explore G [] enum.items (fun a => Or.inl (enum.complete a)) (by simp)

def verdictCode {A : Type u} {G : Graph A} : Verdict G -> Option (List A)
  | .safe _ => none
  | .counterexample <_, p, _> => some p.nodes

theorem solve_safe_iff {A : Type u} (G : Graph A) (enum : Enumeration A) :
    (verdictCode (solve G enum) = none) <-> Safe G := by
  cases solve G enum with
  | safe hs => simp [verdictCode, hs]
  | counterexample p =>
    have hs : not  Safe G := fun h => h p
    obtain <a, p, hb> := p
    simp [verdictCode, hs]

end Cegar
