import Mathlib.Algebra.Order.Ring.Int
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.Ring

/-! Ordered Euclidean ledgers over **all integers**, with arbitrary positive
numerator and denominator. No finite-domain assumption is used here. -/
namespace Omega.Ledger

def step (A B d : ℤ) : ℤ := A*d/B
def remainder (A B d : ℤ) : ℤ := A*d%B

def encode (A B : ℤ) : ℕ → ℤ → ℤ × List ℤ
  | 0, d => (d, [])
  | n+1, d =>
      let tail := encode A B n (step A B d)
      (tail.1, remainder A B d :: tail.2)

def recover (A B t : ℤ) : List ℤ → ℤ
  | [] => t
  | r :: rs => (B*recover A B t rs+r)/A

/-- Admissibility checks every reverse step, not only the final weighted sum. -/
def Admissible (A B t : ℤ) : List ℤ → Prop
  | [] => True
  | r :: rs => Admissible A B t rs ∧ 0 ≤ r ∧ r < B ∧
      (B*recover A B t rs+r)%A=0

theorem euclidean_identity (A B d : ℤ) :
    B*step A B d + remainder A B d = A*d := Int.ediv_add_emod _ _

theorem remainder_bounds (A B d : ℤ) (hB : 0 < B) :
    0 ≤ remainder A B d ∧ remainder A B d < B :=
  ⟨Int.emod_nonneg _ (ne_of_gt hB), Int.emod_lt_of_pos _ hB⟩

theorem one_step_inverse (A B d : ℤ) (hA : A ≠ 0) :
    (B*step A B d+remainder A B d)/A=d := by
  rw [euclidean_identity]
  exact Int.mul_ediv_cancel_left _ hA

theorem roundtrip (A B : ℤ) (hA : A ≠ 0) (n : ℕ) (d : ℤ) :
    recover A B (encode A B n d).1 (encode A B n d).2=d := by
  induction n generalizing d with
  | zero => rfl
  | succ n ih =>
    simp only [encode, recover]
    rw [ih]
    exact one_step_inverse A B d hA

theorem encode_length (A B : ℤ) (n : ℕ) (d : ℤ) :
    (encode A B n d).2.length=n := by
  induction n generalizing d with
  | zero => rfl
  | succ n ih => simp only [encode, List.length_cons, ih]

theorem generated_admissible (A B : ℤ) (hA : A ≠ 0) (hB : 0 < B)
    (n : ℕ) (d : ℤ) :
    Admissible A B (encode A B n d).1 (encode A B n d).2 := by
  induction n generalizing d with
  | zero => trivial
  | succ n ih =>
    simp only [encode, Admissible]
    refine ⟨ih _, (remainder_bounds A B d hB).1,
      (remainder_bounds A B d hB).2, ?_⟩
    rw [roundtrip A B hA, euclidean_identity]
    simp

theorem canonical_division (A B d e r : ℤ) (hB : 0 < B)
    (hr : 0 ≤ r ∧ r < B) (he : A*d=B*e+r) :
    step A B d=e ∧ remainder A B d=r := by
  have hs : step A B d=e := by
    apply (Int.ediv_eq_iff_of_pos hB).mpr
    constructor <;> nlinarith
  refine ⟨hs, ?_⟩
  have hid := euclidean_identity A B d
  rw [hs] at hid
  linarith

theorem admissible_roundtrip (A B : ℤ) (hB : 0 < B) (t : ℤ)
    (rs : List ℤ) (hv : Admissible A B t rs) :
    encode A B rs.length (recover A B t rs)=(t,rs) := by
  induction rs with
  | nil => rfl
  | cons r rs ih =>
    obtain ⟨ht,hr0,hrB,hm⟩ := hv
    have he : A*recover A B t (r::rs)=B*recover A B t rs+r := by
      exact Int.mul_ediv_cancel_of_emod_eq_zero hm
    have hc := canonical_division A B (recover A B t (r::rs))
      (recover A B t rs) r hB ⟨hr0,hrB⟩ he
    simp only [List.length_cons, encode]
    rw [hc.1, hc.2, ih ht]

theorem encoding_injective (A B : ℤ) (hA : A ≠ 0) (n : ℕ) :
    Function.Injective (encode A B n) := by
  intro d e h
  have h' := congrArg (fun z : ℤ × List ℤ => recover A B z.1 z.2) h
  simpa only [roundtrip A B hA] using h'

theorem admissibility_iff (A B : ℤ) (hA : A ≠ 0) (hB : 0 < B)
    (t : ℤ) (rs : List ℤ) :
    Admissible A B t rs ↔ ∃! d, encode A B rs.length d=(t,rs) := by
  constructor
  · intro hv
    refine ⟨recover A B t rs, admissible_roundtrip A B hB t rs hv, ?_⟩
    intro d hd
    have h := roundtrip A B hA rs.length d
    rw [hd] at h
    exact h.symm
  · rintro ⟨d,hd,_⟩
    have h := generated_admissible A B hA hB rs.length d
    simpa only [hd] using h

/-- Position is retained: the head has coefficient A^(length tail), and all
later coefficients acquire a factor B. -/
def value (A B : ℤ) : List ℤ → ℤ
  | [] => 0
  | r::rs => A^rs.length*r+B*value A B rs

theorem weighted_identity (A B : ℤ) (t : ℤ) (rs : List ℤ)
    (hv : Admissible A B t rs) :
    A^rs.length*recover A B t rs=B^rs.length*t+value A B rs := by
  induction rs with
  | nil => simp [recover,value]
  | cons r rs ih =>
    obtain ⟨ht,_,_,hm⟩ := hv
    have he : A*recover A B t (r::rs)=B*recover A B t rs+r :=
      Int.mul_ediv_cancel_of_emod_eq_zero hm
    have hi := ih ht
    simp only [List.length_cons,value,pow_succ]
    calc
      A^rs.length*A*recover A B t (r::rs) =
          A^rs.length*(B*recover A B t rs+r) := by rw [mul_assoc,he]
      _ = B*(A^rs.length*recover A B t rs)+A^rs.length*r := by ring
      _ = B*(B^rs.length*t+value A B rs)+A^rs.length*r := by rw [hi]
      _ = B^rs.length*B*t+(A^rs.length*r+B*value A B rs) := by ring

theorem encoded_weighted_identity (A B : ℤ) (hA : A ≠ 0) (hB : 0 < B)
    (n : ℕ) (d : ℤ) :
    A^n*d=B^n*(encode A B n d).1+value A B (encode A B n d).2 := by
  have hw := weighted_identity A B (encode A B n d).1 (encode A B n d).2
    (generated_admissible A B hA hB n d)
  simpa only [encode_length,roundtrip A B hA] using hw

theorem trajectory_reconstruction (A B : ℤ) (hA : A ≠ 0) (n j : ℕ) (d : ℤ) :
    (step A B)^[j]
      (recover A B (encode A B n d).1 (encode A B n d).2)=(step A B)^[j] d := by
  rw [roundtrip A B hA]

end Omega.Ledger