import Std

/-!
An abstract, kernel-checked flow-lattice theorem. The graph-specific input is
the standard conformal-cycle property: every nonzero edge of an integer
circulation lies on a signed unit cycle conformal to that circulation.

No result about reconfiguration or bounded lifting is assumed by the structure.
The accompanying paper explains the graph instantiation and the independent
standard fact that every modular graph flow has an unbounded integer lift.
-/

namespace FlowReconfiguration

abbrev Vec (E : Type) := E → Int

def add (f g : Vec E) : Vec E := fun e => f e + g e
def sub (f g : Vec E) : Vec E := fun e => f e - g e
def scale (a : Int) (f : Vec E) : Vec E := fun e => a * f e

def Conformal (c f : Vec E) : Prop :=
  ∀ e, (c e = 0) ∨ (c e = 1 ∧ 0 < f e) ∨ (c e = -1 ∧ f e < 0)

/-- Integer circulation lattices equipped with their signed simple cycles. -/
structure FlowSystem (E : Type) where
  edges : List E
  covers : ∀ e, e ∈ edges
  flow : Vec E → Prop
  zero_flow : flow (fun _ => 0)
  add_flow : ∀ {f g}, flow f → flow g → flow (add f g)
  sub_flow : ∀ {f g}, flow f → flow g → flow (sub f g)
  scale_flow : ∀ a {f}, flow f → flow (scale a f)
  cycle : Vec E → Prop
  cycle_flow : ∀ {c}, cycle c → flow c
  cycle_unit : ∀ {c}, cycle c → ∀ e, c e = -1 ∨ c e = 0 ∨ c e = 1
  cycle_neg : ∀ {c}, cycle c → cycle (scale (-1) c)
  conformal_cycle : ∀ {f}, flow f → ∀ e, f e ≠ 0 →
    ∃ c, cycle c ∧ Conformal c f ∧ c e ≠ 0

def Cong (k : Int) (f g : Vec E) : Prop := ∀ e, f e % k = g e % k
def Bounded (k : Int) (f : Vec E) : Prop := ∀ e, -k < f e ∧ f e < k
def Nonzero (f : Vec E) : Prop := ∀ e, f e ≠ 0
def Valid (S : FlowSystem E) (k : Int) (f : Vec E) : Prop :=
  S.flow f ∧ Bounded k f ∧ Nonzero f

def IStep (S : FlowSystem E) (k : Int) (f g : Vec E) : Prop :=
  Valid S k f ∧ Valid S k g ∧ ∃ c, S.cycle c ∧ ∃ a, g = add f (scale a c)

/-- Reconfiguration on residue classes, represented by arbitrary bounded lifts.
The congruence branch identifies different representatives of the same class.
-/
def MStep (S : FlowSystem E) (k : Int) (f g : Vec E) : Prop :=
  Valid S k f ∧ Valid S k g ∧
    (Cong k f g ∨ ∃ c, S.cycle c ∧ ∃ a, Cong k g (add f (scale a c)))

inductive Reach (r : α → α → Prop) : α → α → Prop where
  | refl (a) : Reach r a a
  | step {a b c} : Reach r a b → r b c → Reach r a c

theorem Reach.single (h : r a b) : Reach r a b := .step (.refl a) h

theorem Reach.trans (h₁ : Reach r a b) (h₂ : Reach r b c) : Reach r a c := by
  induction h₂ with
  | refl => exact h₁
  | step _ h ih => exact .step ih h

variable {E : Type} {k : Int} {f g j : Vec E}

theorem Cong.refl (k : Int) (f : Vec E) : Cong k f f := fun _ => rfl
theorem Cong.symm (h : Cong k f g) : Cong k g f := fun e => (h e).symm
theorem Cong.trans (h : Cong k f g) (h' : Cong k g j) : Cong k f j :=
  fun e => (h e).trans (h' e)

theorem Cong.add_right (h : Cong k f g) (j : Vec E) :
    Cong k (add f j) (add g j) := by
  intro e
  change (f e + j e) % k = (g e + j e) % k
  rw [Int.add_emod (f e) (j e) k, Int.add_emod (g e) (j e) k, h e]

theorem Cong.sub_right (h : Cong k f g) (j : Vec E) :
    Cong k (sub f j) (sub g j) := by
  intro e
  change (f e - j e) % k = (g e - j e) % k
  rw [Int.sub_emod (f e) (j e) k, Int.sub_emod (g e) (j e) k, h e]

theorem Cong.sub_multiple (f c : Vec E) (k : Int) :
    Cong k (sub f (scale k c)) f := by
  intro e
  simp [sub, scale, Int.sub_emod]

theorem bounded_nonzero_residue (_hk : 0 < k) (hf : Bounded k f)
    (hn : Nonzero f) : ∀ e, f e % k ≠ 0 := by
  intro e he
  have hb := hf e
  have hz := hn e
  by_cases hpos : 0 ≤ f e
  · have hrem := Int.emod_eq_of_lt hpos hb.2
    omega
  · have hrem := Int.emod_eq_of_lt (show 0 ≤ f e + k from by omega)
      (show f e + k < k from by omega)
    have heq : (f e + k) % k = f e % k := by simp
    omega

theorem shift_safe (hk : 0 < k) {x c : Int}
    (hb : -k < x ∧ x < k) (hc : -1 ≤ c ∧ c ≤ 1)
    (hn : (x+c) % k ≠ 0) : -k < x+c ∧ x+c < k ∧ x+c ≠ 0 := by
  have htop : x+c ≠ k := by
    intro h
    simp [h] at hn
  have hbot : x+c ≠ -k := by
    intro h
    simp [h] at hn
  have hzero : x+c ≠ 0 := by
    intro h
    simp [h] at hn
  omega

def cost (k : Nat) (x : Int) : Nat := x.natAbs - (k-1)
def potential (es : List E) (k : Nat) (f : Vec E) : Nat :=
  (es.map (fun e => cost k (f e))).sum

theorem cost_zero_iff (hk : 0 < k) (x : Int) :
    cost k.toNat x = 0 ↔ -k < x ∧ x < k := by
  have hkn : (k.toNat : Int) = k := Int.toNat_of_nonneg (by omega)
  have habs := Int.natAbs_eq x
  simp only [cost]
  omega

theorem cost_step (k : Nat) (hk : 0 < k) (x c : Int)
    (hc : c = 0 ∨ (c = 1 ∧ 0 < x) ∨ (c = -1 ∧ x < 0)) :
    cost k (x-(k:Int)*c) ≤ cost k x := by
  have h₁ := Int.natAbs_eq x
  have h₂ := Int.natAbs_eq (x-(k:Int)*c)
  simp only [cost]
  rcases hc with h | ⟨h, hx⟩ | ⟨h, hx⟩ <;> subst c <;> simp at h₂ ⊢ <;> omega

theorem cost_step_strict (k : Nat) (hk : 0 < k) (x c : Int)
    (hx : 0 < cost k x) (hn : c ≠ 0)
    (hc : c = 0 ∨ (c = 1 ∧ 0 < x) ∨ (c = -1 ∧ x < 0)) :
    cost k (x-(k:Int)*c) < cost k x := by
  have h₁ := Int.natAbs_eq x
  have h₂ := Int.natAbs_eq (x-(k:Int)*c)
  simp only [cost] at hx ⊢
  rcases hc with h | ⟨h, hx⟩ | ⟨h, hx⟩ <;> subst c <;> simp at h₂ hn ⊢ <;> omega

theorem sum_mono (es : List E) (a b : E → Nat) (h : ∀ e, a e ≤ b e) :
    (es.map a).sum ≤ (es.map b).sum := by
  induction es with
  | nil => simp
  | cons e es ih => simp only [List.map_cons, List.sum_cons]; exact Nat.add_le_add (h e) ih

theorem sum_strict (es : List E) (a b : E → Nat)
    (h : ∀ e, a e ≤ b e) (e : E) (he : e ∈ es) (hlt : a e < b e) :
    (es.map a).sum < (es.map b).sum := by
  induction es with
  | nil => simp at he
  | cons x es ih =>
    simp only [List.map_cons, List.sum_cons]
    have hsum := sum_mono es a b h
    rcases List.mem_cons.mp he with hxe | hem
    · subst x; omega
    · have hs := ih hem; have hx := h x; omega

theorem bounded_lift (S : FlowSystem E) (hk : 0 < k) (f : Vec E) (hf : S.flow f) :
    ∃ g, S.flow g ∧ Bounded k g ∧ Cong k g f := by
  have hkn : (k.toNat : Int) = k := Int.toNat_of_nonneg (by omega)
  have hkp : 0 < k.toNat := by omega
  have aux : ∀ N, ∀ f : Vec E, potential S.edges k.toNat f = N → S.flow f →
      ∃ g, S.flow g ∧ Bounded k g ∧ Cong k g f := by
    intro N
    induction N using Nat.strongRecOn with
    | ind N ih =>
      intro f hN hf
      by_cases hb : Bounded k f
      · exact ⟨f, hf, hb, Cong.refl k f⟩
      · have he : ∃ e, ¬(-k < f e ∧ f e < k) := Classical.not_forall.mp hb
        obtain ⟨e, he⟩ := he
        have hecost : 0 < cost k.toNat (f e) := by
          have hc := cost_zero_iff hk (f e)
          omega
        have hne : f e ≠ 0 := by intro hz; simp [cost, hz] at hecost
        obtain ⟨c, hc, hcf, hce⟩ := S.conformal_cycle hf e hne
        let next := sub f (scale k c)
        have hnext : S.flow next := S.sub_flow hf (S.scale_flow k (S.cycle_flow hc))
        have hlt : potential S.edges k.toNat next < N := by
          rw [← hN]
          apply sum_strict S.edges (fun j => cost k.toNat (next j))
            (fun j => cost k.toNat (f j)) _ e (S.covers e)
          · simpa [next, sub, scale, hkn] using
              cost_step_strict k.toNat hkp (f e) (c e) hecost hce (hcf e)
          · intro j
            simpa [next, sub, scale, hkn] using cost_step k.toNat hkp (f j) (c j) (hcf j)
        obtain ⟨g, hg, hgb, hgc⟩ := ih _ hlt next rfl hnext
        exact ⟨g, hg, hgb, hgc.trans (Cong.sub_multiple f c k)⟩
  exact aux _ f rfl hf

theorem bounded_congruent_difference {x y : Int} (hk : 0 < k)
    (hx : -k < x ∧ x < k) (hy : -k < y ∧ y < k)
    (he : x % k = y % k) : y-x = -k ∨ y-x = 0 ∨ y-x = k := by
  have hrem : (y-x) % k = 0 := by
    rw [Int.sub_emod, he]
    simp
  have hd : y-x = k * ((y-x)/k) := (Int.mul_ediv_cancel_of_emod_eq_zero hrem).symm
  have hdivlo : -2 < (y-x)/k := by
    exact (Int.lt_ediv_iff_mul_lt hk (Int.dvd_of_emod_eq_zero hrem)).mpr (by omega)
  have hdivhi : (y-x)/k < 2 := by
    exact (Int.ediv_lt_iff_lt_mul hk).mpr (by omega)
  have hcases : (y-x)/k = -1 ∨ (y-x)/k = 0 ∨ (y-x)/k = 1 := by omega
  rcases hcases with h | h | h <;> rw [h] at hd <;> simp at hd <;> omega

theorem fiber_coordinate {x y c : Int} (hk : 0 < k)
    (hx : -k < x ∧ x < k) (hy : -k < y ∧ y < k)
    (he : x % k = y % k)
    (hc : c = 0 ∨ (c = 1 ∧ 0 < y-x) ∨ (c = -1 ∧ y-x < 0)) :
    x+k*c = x ∨ x+k*c = y := by
  have hd := bounded_congruent_difference hk hx hy he
  rcases hc with h | ⟨h, hp⟩ | ⟨h, hn⟩ <;> subst c <;> simp <;> omega

theorem fiber_connected (S : FlowSystem E) (hk : 0 < k)
    (hf : Valid S k f) (hg : Valid S k g) (hfg : Cong k f g) :
    Reach (IStep S k) f g := by
  have hkn : (k.toNat : Int) = k := Int.toNat_of_nonneg (by omega)
  have hkp : 0 < k.toNat := by omega
  have aux : ∀ N, ∀ f : Vec E, potential S.edges k.toNat (sub g f) = N →
      Valid S k f → Cong k f g → Reach (IStep S k) f g := by
    intro N
    induction N using Nat.strongRecOn with
    | ind N ih =>
      intro f hN hf hfg
      by_cases heq : f = g
      · subst f; exact .refl g
      · have hex : ∃ e, f e ≠ g e := by
          apply Classical.byContradiction
          intro hn
          have hh : ∀ e, f e = g e := by
            intro e
            apply Classical.byContradiction
            intro he
            exact hn ⟨e, he⟩
          exact heq (funext hh)
        obtain ⟨e, he⟩ := hex
        have hdiff : S.flow (sub g f) := S.sub_flow hg.1 hf.1
        obtain ⟨c, hc, hcf, hce⟩ := S.conformal_cycle hdiff e (by simp [sub]; omega)
        let next := add f (scale k c)
        have hncoord : ∀ j, next j = f j ∨ next j = g j := by
          intro j
          exact fiber_coordinate hk (hf.2.1 j) (hg.2.1 j) (hfg j) (hcf j)
        have hnv : Valid S k next := by
          refine ⟨S.add_flow hf.1 (S.scale_flow k (S.cycle_flow hc)), ?_, ?_⟩
          · intro j
            rcases hncoord j with h | h
            · rw [h]; exact hf.2.1 j
            · rw [h]; exact hg.2.1 j
          · intro j
            rcases hncoord j with h | h
            · rw [h]; exact hf.2.2 j
            · rw [h]; exact hg.2.2 j
        have hng : Cong k next g := by
          intro j
          rcases hncoord j with h | h
          · rw [h]; exact hfg j
          · rw [h]
        have hcost : 0 < cost k.toNat ((sub g f) e) := by
          have hd := bounded_congruent_difference hk (hf.2.1 e) (hg.2.1 e) (hfg e)
          have ha := Int.natAbs_eq (g e - f e)
          simp only [cost, sub]
          omega
        have hid : sub g next = sub (sub g f) (scale k c) := by
          funext j
          simp only [sub, next, add, scale]
          omega
        have hlt : potential S.edges k.toNat (sub g next) < N := by
          rw [← hN, hid]
          apply sum_strict S.edges
            (fun j => cost k.toNat ((sub (sub g f) (scale k c)) j))
            (fun j => cost k.toNat ((sub g f) j)) _ e (S.covers e)
          · simpa [sub, scale, hkn] using
              cost_step_strict k.toNat hkp ((sub g f) e) (c e) hcost hce (hcf e)
          · intro j
            simpa [sub, scale, hkn] using
              cost_step k.toNat hkp ((sub g f) j) (c j) (hcf j)
        have hstep : IStep S k f next := ⟨hf, hnv, c, hc, k, rfl⟩
        exact (Reach.single hstep).trans (ih _ hlt next rfl hnv hng)
  exact aux _ f rfl hf hfg

theorem unit_step_lifts (S : FlowSystem E) (hk : 0 < k)
    (hf : Valid S k f) (hg : Valid S k g) {c : Vec E} (hc : S.cycle c)
    (hcong : Cong k g (add f c)) : Reach (IStep S k) f g := by
  let next := add f c
  have hnres := bounded_nonzero_residue hk hg.2.1 hg.2.2
  have hsafe : ∀ e, -k < next e ∧ next e < k ∧ next e ≠ 0 := by
    intro e
    have hcu := S.cycle_unit hc e
    apply shift_safe hk (hf.2.1 e) (by omega)
    change (add f c e) % k ≠ 0
    rw [← hcong e]
    exact hnres e
  have hv : Valid S k next :=
    ⟨S.add_flow hf.1 (S.cycle_flow hc),
      fun e => ⟨(hsafe e).1, (hsafe e).2.1⟩, fun e => (hsafe e).2.2⟩
  have hs : IStep S k f next := by
    refine ⟨hf, hv, c, hc, 1, ?_⟩
    funext e; simp [next, add, scale]
  exact (Reach.single hs).trans (fiber_connected S hk hv hg hcong.symm)

/-- The new midpoint lemma: modular cycle moves of amplitude two lift. -/
theorem double_step_lifts (S : FlowSystem E) (hk : 0 < k)
    (hf : Valid S k f) (hg : Valid S k g) {c : Vec E} (hc : S.cycle c)
    (hcong : Cong k g (add f (scale 2 c))) : Reach (IStep S k) f g := by
  obtain ⟨h, hh, hb, he⟩ := bounded_lift S hk (add f c)
    (S.add_flow hf.1 (S.cycle_flow hc))
  let left := sub h c
  let right := add h c
  have hleft : Cong k left f := by
    have ht := he.sub_right c
    have hid : sub (add f c) c = f := by funext e; simp [sub, add]
    rw [hid] at ht
    exact ht
  have hright : Cong k right g := by
    have ht := he.add_right c
    have hid : add (add f c) c = add f (scale 2 c) := by
      funext e; simp only [add, scale]; omega
    rw [hid] at ht
    exact ht.trans hcong.symm
  have hfres := bounded_nonzero_residue hk hf.2.1 hf.2.2
  have hgres := bounded_nonzero_residue hk hg.2.1 hg.2.2
  have hlv : Valid S k left := by
    have hs : ∀ e, -k < left e ∧ left e < k ∧ left e ≠ 0 := by
      intro e
      have hcu := S.cycle_unit hc e
      have hres : (h e + -c e) % k ≠ 0 := by
        have hid : h e + -c e = left e := by simp [left, sub, Int.sub_eq_add_neg]
        rw [hid, hleft e]; exact hfres e
      simpa [left, sub, Int.sub_eq_add_neg] using shift_safe hk (hb e) (by omega : -1 ≤ -c e ∧ -c e ≤ 1) hres
    exact ⟨S.sub_flow hh (S.cycle_flow hc),
      fun e => ⟨(hs e).1, (hs e).2.1⟩, fun e => (hs e).2.2⟩
  have hrv : Valid S k right := by
    have hs : ∀ e, -k < right e ∧ right e < k ∧ right e ≠ 0 := by
      intro e
      have hcu := S.cycle_unit hc e
      apply shift_safe hk (hb e) (by omega)
      change right e % k ≠ 0
      rw [hright e]; exact hgres e
    exact ⟨S.add_flow hh (S.cycle_flow hc),
      fun e => ⟨(hs e).1, (hs e).2.1⟩, fun e => (hs e).2.2⟩
  have hstep : IStep S k left right := by
    refine ⟨hlv, hrv, c, hc, 2, ?_⟩
    funext e; simp only [left, right, sub, add, scale]; omega
  exact (fiber_connected S hk hf hlv hleft.symm).trans
    ((Reach.single hstep).trans (fiber_connected S hk hrv hg hright))

theorem small_residue (hk : k = 4 ∨ k = 5) (a : Int) :
    ∃ b : Int, -2 ≤ b ∧ b ≤ 2 ∧ a % k = b % k := by
  rcases hk with rfl | rfl
  · have hlo := Int.emod_nonneg a (by decide : (4 : Int) ≠ 0)
    have hhi := Int.emod_lt_of_pos a (by decide : (0 : Int) < 4)
    by_cases hh : a % 4 ≤ 2
    · exact ⟨a % 4, by omega, hh, by omega⟩
    · exact ⟨a % 4 - 4, by omega, by omega, by omega⟩
  · have hlo := Int.emod_nonneg a (by decide : (5 : Int) ≠ 0)
    have hhi := Int.emod_lt_of_pos a (by decide : (0 : Int) < 5)
    by_cases hh : a % 5 ≤ 2
    · exact ⟨a % 5, by omega, hh, by omega⟩
    · exact ⟨a % 5 - 5, by omega, by omega, by omega⟩

theorem scale_congruence (f c : Vec E) {a b : Int} (h : a % k = b % k) :
    Cong k (add f (scale a c)) (add f (scale b c)) := by
  intro e
  change (f e + a * c e) % k = (f e + b * c e) % k
  rw [Int.add_emod (f e) (a*c e) k, Int.add_emod (f e) (b*c e) k,
      Int.mul_emod a (c e) k, Int.mul_emod b (c e) k, h]

theorem modular_step_lifts (S : FlowSystem E) (hk : k = 4 ∨ k = 5)
    (hs : MStep S k f g) : Reach (IStep S k) f g := by
  have hkp : 0 < k := by omega
  obtain ⟨hf, hg, hcg | ⟨c, hc, a, ha⟩⟩ := hs
  · exact fiber_connected S hkp hf hg hcg
  · obtain ⟨b, hb₁, hb₂, hab⟩ := small_residue hk a
    have hgb := ha.trans (scale_congruence f c hab)
    have hcases : b = -2 ∨ b = -1 ∨ b = 0 ∨ b = 1 ∨ b = 2 := by omega
    rcases hcases with hb | hb | hb | hb | hb <;> subst b
    · have hid : add f (scale (-2) c) = add f (scale 2 (scale (-1) c)) := by
        funext e; simp only [add, scale]; omega
      rw [hid] at hgb
      exact double_step_lifts S hkp hf hg (S.cycle_neg hc) hgb
    · exact unit_step_lifts S hkp hf hg (S.cycle_neg hc) hgb
    · have hid : add f (scale 0 c) = f := by funext e; simp [add, scale]
      rw [hid] at hgb
      exact fiber_connected S hkp hf hg hgb.symm
    · have hid : add f (scale 1 c) = add f c := by funext e; simp [add, scale]
      rw [hid] at hgb
      exact unit_step_lifts S hkp hf hg hc hgb
    · exact double_step_lifts S hkp hf hg hc hgb

theorem integer_step_projects (S : FlowSystem E) (hs : IStep S k f g) :
    MStep S k f g := by
  obtain ⟨hf, hg, c, hc, a, ha⟩ := hs
  refine ⟨hf, hg, Or.inr ⟨c, hc, a, ?_⟩⟩
  rw [ha]
  exact Cong.refl k _

/-- Main theorem: component equivalence for k=4 and k=5 in every finite
conformal signed-unit circulation system. All structural hypotheses are
visible in FlowSystem; bounded lifting and fiber connectivity are proved.
-/
theorem component_equivalence (S : FlowSystem E) (hk : k = 4 ∨ k = 5)
    (f g : Vec E) : Reach (IStep S k) f g ↔ Reach (MStep S k) f g := by
  constructor
  · intro hp
    induction hp with
    | refl => exact .refl _
    | step _ hs ih => exact .step ih (integer_step_projects S hs)
  · intro hp
    induction hp with
    | refl => exact .refl _
    | step _ hs ih => exact ih.trans (modular_step_lifts S hk hs)

theorem four_component_equivalence (S : FlowSystem E) (f g : Vec E) :
    Reach (IStep S 4) f g ↔ Reach (MStep S 4) f g :=
  component_equivalence S (Or.inl rfl) f g

theorem five_component_equivalence (S : FlowSystem E) (f g : Vec E) :
    Reach (IStep S 5) f g ↔ Reach (MStep S 5) f g :=
  component_equivalence S (Or.inr rfl) f g

#print axioms bounded_lift
#print axioms fiber_connected
#print axioms double_step_lifts
#print axioms component_equivalence

end FlowReconfiguration