-- Bank fact: safe-a-wash-cycle-is-a-ring-sell-and-a-ring-buy-each-trade-once (bank/invariants.json)
-- Machine checked with Lean 4.22.0, core only, no Mathlib (bank/proofs/lean-toolchain).
-- Check it: lean SafeWash.lean   (the #print axioms lines list only Lean's standard axioms)
/-
HARMONIC ASSAY (first named SAFE): a wash cycle is a ring sell and a later ring buy within fifty blocks
and a quarter of the size, and no trade is counted in two cycles.

The model is the function the agent runs (agent/src/safe.js, washCycles):
the trades are filtered to the ring and kept in chain order; each sell not
yet used, of positive size, takes the first later trade, before the first
one more than fifty blocks on, that is an unused buy whose size is within a
quarter of the larger of the two; both are marked used.
-/

structure Trade where
  sell : Bool
  inRing : Bool
  size : Nat
  block : Nat
deriving Inhabited

def window : Nat := 50

/-- sizes within a quarter of the larger: (hi - lo) / hi ≤ 1/4, in integers -/
def tol (a b : Nat) : Bool := decide (4 * (max a b - min a b) ≤ max a b)

def at' (L : List Trade) (i : Nat) : Trade := L.getD i default

/-- the trades after i, in order, up to the first more than `window` blocks on -/
def cands (L : List Trade) (i : Nat) : List Nat :=
  ((List.range L.length).filter (fun j => decide (i < j))).takeWhile
    (fun j => decide ((at' L j).block ≤ (at' L i).block + window))

def fits (L : List Trade) (i : Nat) (used : List Nat) (j : Nat) : Bool :=
  !(used.contains j) && !(at' L j).sell && tol (at' L i).size (at' L j).size

def scan (L : List Trade) (i : Nat) (used : List Nat) : Option Nat :=
  (cands L i).find? (fits L i used)

def step (L : List Trade) (st : List Nat × List (Nat × Nat)) (i : Nat) :
    List Nat × List (Nat × Nat) :=
  if (at' L i).sell && decide (0 < (at' L i).size) && !(st.1.contains i) then
    match scan L i st.1 with
    | some j => (j :: i :: st.1, st.2 ++ [(i, j)])
    | none => st
  else st

def pairsOf (L : List Trade) : List (Nat × Nat) :=
  ((List.range L.length).foldl (step L) ([], [])).2

/-- the agent's washCycles: the ring's trades, paired -/
def washCycles (T : List Trade) : List (Nat × Nat) := pairsOf (T.filter (·.inRing))

def flat (ps : List (Nat × Nat)) : List Nat := ps.flatMap (fun p => [p.1, p.2])

/-- what one cycle must be -/
def good (L : List Trade) (p : Nat × Nat) : Prop :=
  p.1 < p.2 ∧ p.2 < L.length ∧ (at' L p.1).sell = true ∧ (at' L p.2).sell = false ∧
  0 < (at' L p.1).size ∧
  (at' L p.2).block ≤ (at' L p.1).block + window ∧
  tol (at' L p.1).size (at' L p.2).size = true

structure WashInv (L : List Trade) (st : List Nat × List (Nat × Nat)) : Prop where
  good : ∀ p ∈ st.2, good L p
  used : ∀ k, k ∈ st.1 ↔ k ∈ flat st.2
  nodup : (flat st.2).Nodup

theorem flat_append (ps : List (Nat × Nat)) (i j : Nat) : flat (ps ++ [(i, j)]) = flat ps ++ [i, j] := by
  simp [flat, List.flatMap_append]

theorem scan_spec (L : List Trade) (i : Nat) (used : List Nat) (j : Nat)
    (h : scan L i used = some j) :
    i < j ∧ j < L.length ∧ j ∉ used ∧ (at' L j).sell = false ∧
    (at' L j).block ≤ (at' L i).block + window ∧ tol (at' L i).size (at' L j).size = true := by
  unfold scan at h
  have hp := List.find?_some h
  have hm := List.mem_of_find?_eq_some h
  have hw := (List.all_eq_true.mp (List.all_takeWhile (l := (List.range L.length).filter (fun j => decide (i < j))))) j hm
  have hf := (List.takeWhile_sublist _).subset hm
  simp only [List.mem_filter, List.mem_range, decide_eq_true_eq] at hf
  simp only [fits, Bool.and_eq_true, Bool.not_eq_true', List.contains_eq_mem,
    decide_eq_false_iff_not] at hp
  simp only [decide_eq_true_eq] at hw
  exact ⟨hf.2, hf.1, hp.1.1, hp.1.2, hw, hp.2⟩

theorem step_inv (L : List Trade) (st : List Nat × List (Nat × Nat)) (i : Nat)
    (hI : WashInv L st) : WashInv L (step L st i) := by
  unfold step
  split
  · rename_i hc
    simp only [Bool.and_eq_true, Bool.not_eq_true', decide_eq_true_eq, List.contains_eq_mem,
      decide_eq_false_iff_not] at hc
    obtain ⟨⟨hsell, hpos⟩, hiu⟩ := hc
    split
    · rename_i j hj
      obtain ⟨hij, hjn, hju, hjb, hjw, hjt⟩ := scan_spec L i st.1 j hj
      have hif : i ∉ flat st.2 := fun h => hiu ((hI.used i).2 h)
      have hjf : j ∉ flat st.2 := fun h => hju ((hI.used j).2 h)
      refine ⟨?_, ?_, ?_⟩
      · intro p hp
        simp only [List.mem_append, List.mem_singleton] at hp
        rcases hp with hp | rfl
        · exact hI.good p hp
        · exact ⟨hij, hjn, hsell, hjb, hpos, hjw, hjt⟩
      · intro k
        simp only [flat_append, List.mem_cons, List.mem_append,
          List.not_mem_nil, or_false]
        rw [hI.used k]
        constructor
        · rintro (h | h | h) <;> simp_all
        · rintro (h | h | h) <;> simp_all
      · rw [flat_append, List.nodup_append]
        refine ⟨hI.nodup, ?_, ?_⟩
        · simp only [List.nodup_cons, List.mem_singleton, List.not_mem_nil, not_false_eq_true,
            List.nodup_nil, and_true]
          omega
        · intro a ha b hb
          simp only [List.mem_cons, List.not_mem_nil, or_false] at hb
          rcases hb with rfl | rfl
          · exact fun e => hif (e ▸ ha)
          · exact fun e => hjf (e ▸ ha)
    · exact hI
  · exact hI

theorem fold_inv (L : List Trade) : ∀ (xs : List Nat) (st : List Nat × List (Nat × Nat)),
    WashInv L st → WashInv L (xs.foldl (step L) st)
  | [], _, h => h
  | x :: xs, st, h => fold_inv L xs (step L st x) (step_inv L st x h)

theorem inv_start (L : List Trade) : WashInv L ([], []) :=
  ⟨by simp, by simp [flat], by simp [flat]⟩

/-- every cycle is a ring sell then a later ring buy inside the window and the tolerance -/
theorem wash_cycle_sound (T : List Trade) :
    ∀ p ∈ washCycles T,
      let L := T.filter (·.inRing)
      good L p ∧ (at' L p.1).inRing = true ∧ (at' L p.2).inRing = true := by
  intro p hp L
  have hg := (fold_inv L (List.range L.length) ([], []) (inv_start L)).good p hp
  have mem : ∀ k, k < L.length → at' L k ∈ L := by
    intro k hk
    simp only [at', List.getD_eq_getElem?_getD, List.getElem?_eq_getElem hk, Option.getD_some]
    exact List.getElem_mem hk
  have ring : ∀ k, k < L.length → (at' L k).inRing = true := by
    intro k hk
    have := mem k hk
    simp only [L, List.mem_filter] at this
    exact this.2
  exact ⟨hg, ring _ (Nat.lt_trans hg.1 hg.2.1), ring _ hg.2.1⟩

/-- no trade is counted in two cycles -/
theorem wash_cycle_each_trade_once (T : List Trade) : (flat (washCycles T)).Nodup :=
  (fold_inv _ _ ([], []) (inv_start _)).nodup

/-- distinct elements, all drawn from l₂, are no more than l₂ has -/
theorem length_le_of_nodup_subset : ∀ (l₁ l₂ : List Nat), l₁.Nodup → (∀ x ∈ l₁, x ∈ l₂) →
    l₁.length ≤ l₂.length
  | [], _, _, _ => Nat.zero_le _
  | a :: as, l₂, hn, hs => by
    rw [List.nodup_cons] at hn
    have ha : a ∈ l₂ := hs a (List.mem_cons_self ..)
    have hsub : ∀ x ∈ as, x ∈ l₂.erase a := by
      intro x hx
      have hne : x ≠ a := fun e => hn.1 (e ▸ hx)
      exact (List.mem_erase_of_ne hne).2 (hs x (List.mem_cons_of_mem _ hx))
    have ih := length_le_of_nodup_subset as (l₂.erase a) hn.2 hsub
    rw [List.length_erase_of_mem ha] at ih
    have := List.length_pos_of_mem ha
    simp only [List.length_cons]
    omega

theorem snd_sublist_flat : ∀ (ps : List (Nat × Nat)), (ps.map Prod.snd).Sublist (flat ps)
  | [] => by simp [flat]
  | p :: ps => by
    simp only [List.map_cons, flat, List.flatMap_cons, List.cons_append, List.nil_append]
    exact List.Sublist.cons _ (List.Sublist.cons₂ _ (snd_sublist_flat ps))

theorem fst_sublist_flat : ∀ (ps : List (Nat × Nat)), (ps.map Prod.fst).Sublist (flat ps)
  | [] => by simp [flat]
  | p :: ps => by
    simp only [List.map_cons, flat, List.flatMap_cons, List.cons_append, List.nil_append]
    exact List.Sublist.cons₂ _ (List.Sublist.cons _ (fst_sublist_flat ps))

/-- the number of cycles is at most the ring's sells and at most the ring's buys -/
theorem wash_cycles_bounded (T : List Trade) :
    let L := T.filter (·.inRing)
    (washCycles T).length ≤ ((List.range L.length).filter (fun k => (at' L k).sell)).length ∧
    (washCycles T).length ≤ ((List.range L.length).filter (fun k => !(at' L k).sell)).length := by
  intro L
  have hI := fold_inv L (List.range L.length) ([], []) (inv_start L)
  have hnd : (flat (washCycles T)).Nodup := hI.nodup
  have hg : ∀ p ∈ washCycles T, good L p := hI.good
  constructor
  · have h1 := length_le_of_nodup_subset ((washCycles T).map Prod.fst)
      ((List.range L.length).filter (fun k => (at' L k).sell))
      ((fst_sublist_flat _).nodup hnd) (by
        intro x hx
        obtain ⟨p, hp, rfl⟩ := List.mem_map.1 hx
        have g := hg p hp
        simp only [List.mem_filter, List.mem_range]
        exact ⟨Nat.lt_trans g.1 g.2.1, g.2.2.1⟩)
    rw [List.length_map] at h1; exact h1
  · have h2 := length_le_of_nodup_subset ((washCycles T).map Prod.snd)
      ((List.range L.length).filter (fun k => !(at' L k).sell))
      ((snd_sublist_flat _).nodup hnd) (by
        intro x hx
        obtain ⟨p, hp, rfl⟩ := List.mem_map.1 hx
        have g := hg p hp
        simp only [List.mem_filter, List.mem_range, Bool.not_eq_true']
        exact ⟨g.2.1, g.2.2.2.1⟩)
    rw [List.length_map] at h2; exact h2

#print axioms wash_cycle_sound
#print axioms wash_cycle_each_trade_once
#print axioms wash_cycles_bounded