import Mathlib.Algebra.BigOperators.Field
import Mathlib.Algebra.BigOperators.Group.List.Basic
import Mathlib.Tactic.FieldSimp
import Mathlib.Tactic.Ring

/-!
Richardson extrapolation: exact cancellation at arbitrary finite order.
Historical source: L. F. Richardson (1911), DOI 10.1098/rsta.1911.0009.
This is a modern algebraic reconstruction on a geometrically refined grid,
not a formalization of the paper's differential equations or numerical errors.
-/
namespace PaperToLean.Richardson
open scoped BigOperators
variable {K : Type*} [Field K]

def step (q : K) (s : ℕ → K) (n : ℕ) : K :=
  (s (n + 1) - q * s n) / (1 - q)

def extrapolate : List K → (ℕ → K) → ℕ → K
  | [], s => s
  | q :: qs, s => step q (extrapolate qs s)

def response (qs : List K) (r : K) : K :=
  (qs.map (fun q => (r - q) / (1 - q))).prod

theorem step_add (q : K) (s t : ℕ → K) :
    step q (fun n => s n + t n) = fun n => step q s n + step q t n := by
  funext n
  simp only [step, div_eq_mul_inv]
  ring

theorem extrapolate_add (qs : List K) (s t : ℕ → K) :
    extrapolate qs (fun n => s n + t n) =
      fun n => extrapolate qs s n + extrapolate qs t n := by
  induction qs with
  | nil => rfl
  | cons q qs ih => simp only [extrapolate, ih, step_add]

theorem extrapolate_sum {ι : Type*} (qs : List K) (I : Finset ι)
    (f : ι → ℕ → K) :
    extrapolate qs (fun n => ∑ i ∈ I, f i n) =
      fun n => ∑ i ∈ I, extrapolate qs (f i) n := by
  classical
  induction I using Finset.induction_on with
  | empty =>
    simp only [Finset.sum_empty]
    induction qs with
    | nil => rfl
    | cons q qs ih =>
      funext n
      simp [extrapolate, ih, step]
  | @insert i I hi ih =>
    simp only [Finset.sum_insert hi, extrapolate_add, ih]

theorem geometric_response (qs : List K) (a r : K) :
    extrapolate qs (fun n => a * r ^ n) =
      fun n => a * r ^ n * response qs r := by
  induction qs with
  | nil => simp [extrapolate, response]
  | cons q qs ih =>
    funext n
    simp only [extrapolate, ih, step, response, List.map_cons, List.prod_cons, pow_succ]
    ring

theorem response_vanishes (qs : List K) (r : K) (hr : r ∈ qs) :
    response qs r = 0 := by
  apply List.prod_eq_zero_iff.mpr
  exact List.mem_map.mpr ⟨r, hr, by simp⟩

theorem constant_preserved (qs : List K) (L : K)
    (h : ∀ q ∈ qs, q ≠ 1) : extrapolate qs (fun _ => L) = fun _ => L := by
  induction qs with
  | nil => rfl
  | cons q qs ih =>
    have hq : 1 - q ≠ 0 := sub_ne_zero.mpr (Ne.symm (h q (by simp)))
    rw [extrapolate, ih (fun r hr => h r (by simp [hr]))]
    funext n
    dsimp [step]
    field_simp

/-- Any finite mixture of the listed modes is eliminated, at every index. -/
theorem finite_spectrum_exact {ι : Type*} (I : Finset ι) (qs : List K)
    (L : K) (a r : ι → K) (h : ∀ q ∈ qs, q ≠ 1)
    (hr : ∀ i ∈ I, r i ∈ qs) (n : ℕ) :
    extrapolate qs (fun n => L + ∑ i ∈ I, a i * r i ^ n) n = L := by
  rw [extrapolate_add, constant_preserved qs L h, extrapolate_sum]
  have hz : (∑ i ∈ I, extrapolate qs (fun n => a i * r i ^ n) n) = 0 := by
    apply Finset.sum_eq_zero
    intro i hi
    simp only [geometric_response, response_vanishes qs (r i) (hr i hi), mul_zero]
  simp only [hz, add_zero]

/-- A polynomial grid error becomes a finite spectrum under h_n = h * t^n. -/
theorem grid_power (h t : K) (n p : ℕ) :
    (h * t ^ n) ^ p = h ^ p * (t ^ p) ^ n := by
  rw [mul_pow, ← pow_mul, Nat.mul_comm n p, pow_mul]

/-- Exactness for any finite polynomial error expansion on a geometric grid. -/
theorem polynomial_grid_exact {ι : Type*} (I : Finset ι) (qs : List K)
    (L h t : K) (a : ι → K) (p : ι → ℕ) (hq : ∀ q ∈ qs, q ≠ 1)
    (hp : ∀ i ∈ I, t ^ p i ∈ qs) (n : ℕ) :
    extrapolate qs (fun n => L + ∑ i ∈ I, a i * (h * t ^ n) ^ p i) n = L := by
  simp_rw [grid_power, ← mul_assoc]
  exact finite_spectrum_exact I qs L (fun i => a i * h ^ p i)
    (fun i => t ^ p i) hq hp n

#print axioms step_add
#print axioms extrapolate_add
#print axioms extrapolate_sum
#print axioms geometric_response
#print axioms response_vanishes
#print axioms constant_preserved
#print axioms finite_spectrum_exact
#print axioms grid_power
#print axioms polynomial_grid_exact
end PaperToLean.Richardson