import Cegar.Core

namespace Cegar

universe u v

section May

variable {S : Type u} {A : Type v} [DecidableEq A]

def may (G : Graph S) (states : Enumeration S) (q : S -> A) : Graph A where
  initial a := states.items.any fun x => decide (q x = a) && G.initial x
  bad a := states.items.any fun x => decide (q x = a) && G.bad x
  edge a b := states.items.any fun x => states.items.any fun y =>
    decide (q x = a) && decide (q y = b) && G.edge x y

variable {G : Graph S} {states : Enumeration S} {q : S -> A}

theorem may_initial {a : A} :
    (may G states q).initial a = true <-> exists  x, q x = a /\ G.initial x = true := by
  simp [may, states.complete]

theorem may_bad {a : A} :
    (may G states q).bad a = true <-> exists  x, q x = a /\ G.bad x = true := by
  simp [may, states.complete]

theorem may_edge {a b : A} :
    (may G states q).edge a b = true <->
      exists  x y, q x = a /\ q y = b /\ G.edge x y = true := by
  simp [may, states.complete, and_assoc]

def mapWalk {x : S} : Walk G x -> Walk (may G states q) (q x)
  | .start x h => .start (q x) (may_initial.mpr <x, rfl, h>)
  | .step p y h => .step (mapWalk p) (q y) (may_edge.mpr <_, y, rfl, rfl, h>)

theorem safe_of_may_safe (h : Safe (may G states q)) : Safe G :=
  fun <x, p, hb> => h <q x, mapWalk p, may_bad.mpr <x, rfl, hb>>

inductive Lift (G : Graph S) (states : Enumeration S) (q : S -> A) :
    {a : A} -> Walk (may G states q) a -> S -> Type (max u v)
  | start {a : A} {h : (may G states q).initial a = true} {x : S} :
    q x = a -> G.initial x = true -> Lift G states q (.start a h) x
  | step {a b : A} {w : Walk (may G states q) a}
      {h : (may G states q).edge a b = true} {x y : S} :
    Lift G states q w x -> q y = b -> G.edge x y = true ->
      Lift G states q (.step w b h) y

def Lift.walk {a : A} {w : Walk (may G states q) a} {x : S} :
    Lift G states q w x -> Walk G x
  | .start _ h => .start _ h
  | .step p _ h => .step p.walk _ h

theorem Lift.endpoint {a : A} {w : Walk (may G states q) a} {x : S}
    (l : Lift G states q w x) : q x = a := by
  cases l with
  | start e _ => exact e
  | step _ e _ => exact e

theorem Lift.nodes {a : A} {w : Walk (may G states q) a} {x : S}
    (l : Lift G states q w x) : l.walk.nodes.map q = w.nodes := by
  induction l with
  | start e _ => simp [Lift.walk, Walk.nodes, e]
  | step _ e _ ih => simp [Lift.walk, Walk.nodes, e, ih]

theorem lift_of_nodes {a : A} {x : S}
    (w : Walk (may G states q) a) (p : Walk G x)
    (same : p.nodes.map q = w.nodes) : Nonempty (Lift G states q w x) := by
  induction w generalizing x with
  | start a hi =>
    cases p with
    | start x hx =>
      have h : q x = a := by simpa [Walk.nodes] using same
      exact <.start h hx>
    | step previous x hx =>
      have h : previous.nodes.map q ++ [q x] = [] ++ [a] := by
        simpa [Walk.nodes] using same
      have empty := List.append_inj_left' h (by simp)
      simpa using (show (previous.nodes.map q).length != 0 from by
        intro he
        have hz : previous.nodes = [] := List.length_eq_zero_iff.mp (by simpa using he)
        have hm := previous.last_mem
        rw [hz] at hm
        contradiction) (congrArg List.length empty)
  | step w b he ih =>
    cases p with
    | start x hx =>
      have h : [] ++ [q x] = w.nodes ++ [b] := by
        simpa [Walk.nodes] using same
      have empty := List.append_inj_left' h (by simp)
      have hm := w.last_mem
      rw [<- empty] at hm
      contradiction
    | step previous x hx =>
      have h : previous.nodes.map q ++ [q x] = w.nodes ++ [b] := by
        simpa [Walk.nodes] using same
      obtain <hl, hr> := List.append_inj' h (by simp)
      obtain <l> := ih previous hl
      exact <.step l (by simpa using hr) hx>

def LiftedBad {a : A} (w : Walk (may G states q) a) :=
  Sigma  x, { _l : Lift G states q w x // G.bad x = true }

def LiftedBad.walk {a : A} {w : Walk (may G states q) a}
    (l : LiftedBad w) : BadWalk G :=
  <l.1, l.2.1.walk, l.2.2>

structure Table {a : A} (w : Walk (may G states q) a) where
  rows : List (Sigma  x, Lift G states q w x)
  complete : forall  x, Lift G states q w x -> exists  r  in  rows, r.1 = x

def extendRow {a b : A} {w : Walk (may G states q) a}
    (h : (may G states q).edge a b = true) (old : Table w) (y : S) :
    Option (Sigma  x, Lift G states q (.step w b h) x) :=
  if hy : q y = b then
    match pick (fun (r : Sigma  x, Lift G states q w x) => if he : G.edge r.1 y = true then
        some (show Lift G states q (.step w b h) y from .step r.2 hy he)
        else none) old.rows with
    | .found _ p _ _ _ _ => some <y, p>
    | .missing _ => none
  else none

theorem extendRow_complete {a b : A} {w : Walk (may G states q) a}
    (h : (may G states q).edge a b = true) (old : Table w) (y : S)
    (l : Lift G states q (.step w b h) y) :
    exists  r, extendRow h old y = some r /\ r.1 = y := by
  cases l with
  | step p hy he =>
    obtain <r, hr, rfl> := old.complete _ p
    unfold extendRow
    simp only [hy, dif_pos]
    split
    . exact <_, rfl, rfl>
    . rename_i miss _
      have contradictionAt := miss r hr
      simp [he] at contradictionAt

def startRow {a : A} (h : (may G states q).initial a = true) (x : S) :
    Option (Sigma  y, Lift G states q (.start a h) y) :=
  if hq : q x = a then
    if hi : G.initial x = true then some <x, .start hq hi> else none
  else none

def tabulate {a : A} (w : Walk (may G states q) a) : Table w :=
  match w with
  | .start _ h =>
    <states.items.filterMap (startRow h), by
      intro x l
      cases l with
      | start hq hi =>
        refine <<x, .start hq hi>, List.mem_filterMap.mpr <x, states.complete x, ?_>, rfl>
        simp [startRow, hq, hi]>
  | .step w _ h =>
    let old := tabulate w
    <states.items.filterMap (extendRow h old), by
      intro x l
      obtain <r, hr, hx> := extendRow_complete h old x l
      exact <r, List.mem_filterMap.mpr <x, states.complete x, hr>, hx>>

def checkBad {a : A} (w : Walk (may G states q) a) : Checked (LiftedBad w) :=
  let table := tabulate w
  let candidate (r : Sigma  x, Lift G states q w x) : Option (LiftedBad w) :=
    if hb : G.bad r.1 = true then some <r.1, r.2, hb> else none
  match pick candidate table.rows with
  | .found _ l _ _ _ _ => .yes l
  | .missing miss => .no (fun <x, l, hb> => by
      obtain <r, hr, rfl> := table.complete x l
      have contradictionAt := miss r hr
      simp [candidate, hb] at contradictionAt)

def UniqueAlong {a : A} (w : Walk (may G states q) a) : Prop :=
  forall  x y, q x = q y -> q x  in  w.nodes -> x = y

theorem lift_of_unique {a : A} (w : Walk (may G states q) a)
    (unique : UniqueAlong w) : Nonempty (Sigma  x, Lift G states q w x) := by
  induction w with
  | start a hi =>
    obtain <x, hx, h> := may_initial.mp hi
    exact <<x, .start hx h>>
  | @step a w b he ih =>
    have up : UniqueAlong w := fun x y eq hx =>
      unique x y eq (List.mem_append_left _ hx)
    obtain <z, l> := ih up
    obtain <x, y, hx, hy, h> := may_edge.mp he
    have e : z = x := unique z x (l.endpoint.trans hx.symm)
      (l.endpoint.symm  |>  List.mem_append_left [b] w.last_mem)
    exact <<y, .step l hy (e.symm  |>  h)>>

theorem liftBad_of_unique {a : A} {w : Walk (may G states q) a}
    (bad : (may G states q).bad a = true) (unique : UniqueAlong w) :
    Nonempty (LiftedBad w) := by
  obtain <x, l> := lift_of_unique w unique
  obtain <y, hy, hb> := may_bad.mp bad
  have e : x = y := unique x y (l.endpoint.trans hy.symm) (l.endpoint.symm  |>  w.last_mem)
  exact <<x, l, e.symm  |>  hb>>

end May

section Loop

variable {S : Type u} {B : Type v} [DecidableEq S] [DecidableEq B]

local instance : DecidableEq (Sum S B) :=
  fun x y => match x, y with
    | .inl a, .inl b => if h : a = b then isTrue (congrArg Sum.inl h)
      else isFalse (fun e => h (Sum.inl.inj e))
    | .inr a, .inr b => if h : a = b then isTrue (congrArg Sum.inr h)
      else isFalse (fun e => h (Sum.inr.inj e))
    | .inl _, .inr _ => isFalse (fun e => nomatch e)
    | .inr _, .inl _ => isFalse (fun e => nomatch e)

def key (seed : S -> B) (remaining : List S) (x : S) : Sum S B :=
  if x  in  remaining then .inr (seed x) else .inl x

def visits (path : List (Sum S B)) (k : Sum S B) : Bool :=
  path.any fun a => decide (a = k)

@[simp] theorem visits_iff (path : List (Sum S B)) (k : Sum S B) :
    visits path k = true <-> k  in  path := by
  simp [visits]

def cut (seed : S -> B) (remaining : List S) (path : List (Sum S B)) : List S :=
  remaining.filter fun x => !(visits path (key seed remaining x))

theorem cut_sub {seed : S -> B} {remaining : List S} {path : List (Sum S B)}
    {x : S} (h : x  in  cut seed remaining path) :
    x  in  remaining /\ key seed remaining x  notin  path := by
  simpa [cut, List.mem_filter, Bool.eq_false_iff] using h

def coarsen (seed : S -> B) (remaining : List S) : Sum S B -> Sum S B
  | .inl x => key seed remaining x
  | .inr b => .inr b

theorem cut_refines (seed : S -> B) (remaining : List S) (path : List (Sum S B)) (x : S) :
    coarsen seed remaining (key seed (cut seed remaining path) x) =
      key seed remaining x := by
  by_cases h : x  in  cut seed remaining path
  . simp [key, coarsen, h, (cut_sub h).1]
  . simp [key, coarsen, h]

theorem cut_unique (seed : S -> B) (remaining : List S) (path : List (Sum S B))
    (x y : S) (same : key seed (cut seed remaining path) x =
      key seed (cut seed remaining path) y)
    (visit : key seed remaining x  in  path) : x = y := by
  have hx : x  notin  cut seed remaining path := fun h => (cut_sub h).2 visit
  rw [key, if_neg hx] at same
  by_cases hy : y  in  cut seed remaining path
  . simp [key, hy] at same
  . simpa [key, hy] using same

theorem old_unique_of_no_cut {G : Graph S} {states : Enumeration S}
    {seed : S -> B} {remaining : List S} {a : Sum S B}
    {w : Walk (may G states (key seed remaining)) a}
    (noCut : forall  x  in  remaining, key seed remaining x  notin  w.nodes) : UniqueAlong w := by
  intro x y same visit
  have hx : x  notin  remaining := fun h => noCut x h visit
  rw [key, if_neg hx] at same
  by_cases hy : y  in  remaining
  . simp [key, hy] at same
  . simpa [key, hy] using same

theorem cut_shorter {G : Graph S} {states : Enumeration S}
    {seed : S -> B} {remaining : List S} {a : Sum S B}
    {w : Walk (may G states (key seed remaining)) a}
    (bad : (may G states (key seed remaining)).bad a = true)
    (refute : LiftedBad w -> False) :
    (cut seed remaining w.nodes).length < remaining.length := by
  by_cases h : remaining.any (fun x => visits w.nodes (key seed remaining x)) = true
  . obtain <x, hx, visit> := List.any_eq_true.mp h
    apply List.length_filter_lt_length_iff_exists.mpr
    exact <x, hx, by simp [visit]>
  . have noCut : forall  x  in  remaining, key seed remaining x  notin  w.nodes :=
      fun x hx hv => h (List.any_eq_true.mpr <x, hx, (visits_iff _ _).2 hv>)
    obtain <l> := liftBad_of_unique bad (old_unique_of_no_cut noCut)
    exact (refute l).elim

theorem cut_eliminates {G : Graph S} {states : Enumeration S}
    {seed : S -> B} {remaining : List S} {a : Sum S B}
    (w : Walk (may G states (key seed remaining)) a)
    (refute : LiftedBad w -> False)
    {b : Sum S B}
    (fine : Walk (may G states (key seed (cut seed remaining w.nodes))) b)
    (bad : (may G states (key seed (cut seed remaining w.nodes))).bad b = true)
    (projects : fine.nodes.map (coarsen seed remaining) = w.nodes) : False := by
  have unique : UniqueAlong fine := by
    intro x y same visit
    apply cut_unique seed remaining w.nodes x y same
    let oldKey := coarsen seed remaining (key seed (cut seed remaining w.nodes) x)
    have hv : oldKey  in  fine.nodes.map (coarsen seed remaining) :=
      List.mem_map.mpr <_, visit, rfl>
    have hw : oldKey  in  w.nodes :=
      (congrArg (fun nodes => oldKey  in  nodes) projects).mp hv
    simpa only [oldKey, cut_refines] using hw
  obtain <<x, l, hx>> := liftBad_of_unique bad unique
  have matched : l.walk.nodes.map (key seed remaining) = w.nodes := by
    let c := coarsen seed remaining
    let q' := key seed (cut seed remaining w.nodes)
    calc
      l.walk.nodes.map (key seed remaining) = l.walk.nodes.map (c  o  q') :=
        List.map_congr_left (fun y _ => (cut_refines seed remaining w.nodes y).symm)
      _ = (l.walk.nodes.map q').map c := List.map_map.symm
      _ = fine.nodes.map c := congrArg (List.map c) l.nodes
      _ = w.nodes := projects
  obtain <old> := lift_of_nodes w l.walk matched
  exact refute <x, old, hx>

structure Run {S : Type u} (G : Graph S) where
  verdict : Verdict G
  rounds : Nat

def runFrom (G : Graph S) (states : Enumeration S)
    (seeds : Enumeration B) (seed : S -> B) (remaining : List S) : Run G :=
  match solve (may G states (key seed remaining)) (sumEnumeration states seeds) with
  | .safe h => <.safe (safe_of_may_safe h), 1>
  | .counterexample <a, p, _bad> =>
    match checkBad p with
    | .yes l => <.counterexample l.walk, 1>
    | .no refute =>
      let next := runFrom G states seeds seed (cut seed remaining p.nodes)
      <next.verdict, next.rounds + 1>
termination_by remaining.length
decreasing_by exact cut_shorter _bad refute

def run (G : Graph S) (states : Enumeration S) (seeds : Enumeration B) (seed : S -> B) :
    Run G :=
  runFrom G states seeds seed states.items

theorem run_safe_iff (G : Graph S) (states : Enumeration S) (seeds : Enumeration B)
    (seed : S -> B) : verdictCode (run G states seeds seed).verdict = none <-> Safe G := by
  cases (run G states seeds seed).verdict with
  | safe h => simp [verdictCode, h]
  | counterexample p =>
    have hs : not  Safe G := fun h => h p
    obtain <a, p, hb> := p
    simp [verdictCode, hs]

end Loop

end Cegar
