import Mathlib.Algebra.Order.BigOperators.Group.Finset
import Mathlib.Algebra.BigOperators.Ring.Finset
import Mathlib.Data.Real.Basic
import Mathlib.Tactic.FieldSimp
import Mathlib.Tactic.Ring
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.Positivity

/-!
Selected results for Kaczmarz's 1937 row-projection method.
Vectors have arbitrary finite dimension. sqdist is squared Euclidean distance,
not the default supremum norm on a function space. The relaxed step extends
the original unrelaxed update; no convergence rate is assumed or advertised.
-/
namespace PaperToLean.Kaczmarz
noncomputable section
open scoped BigOperators
variable {n : ℕ}
abbrev Vec (n : ℕ) := Fin n → ℝ

def dot (x y : Vec n) : ℝ := ∑ i, x i * y i
def sqdist (x y : Vec n) : ℝ := ∑ i, (x i - y i) ^ 2
def step (a : Vec n) (b ω : ℝ) (x : Vec n) : Vec n :=
  fun i => x i + ω * (b - dot a x) / dot a a * a i

theorem dot_self_nonneg (a : Vec n) : 0 ≤ dot a a := by
  exact Finset.sum_nonneg (fun i _ => mul_self_nonneg (a i))

theorem dot_self_eq_zero (a : Vec n) : dot a a = 0 ↔ a = 0 := by
  constructor
  · intro h
    have hzero := (Finset.sum_eq_zero_iff_of_nonneg
      (fun i (_ : i ∈ Finset.univ) => mul_self_nonneg (a i))).mp h
    funext i
    exact mul_self_eq_zero.mp (hzero i (Finset.mem_univ i))
  · rintro rfl
    simp [dot]

theorem row_residual (a x : Vec n) (b ω : ℝ) (ha : dot a a ≠ 0) :
    dot a (step a b ω x) - b = (1 - ω) * (dot a x - b) := by
  have hs (c : ℝ) : dot a (fun i => x i + c * a i) = dot a x + c * dot a a := by
    simp only [dot, mul_add, Finset.sum_add_distrib, Finset.mul_sum]
    congr 1
    apply Finset.sum_congr rfl
    intro i _
    ring
  unfold step
  rw [hs]
  field_simp
  ring

theorem satisfies_row (a x : Vec n) (b : ℝ) (ha : dot a a ≠ 0) :
    dot a (step a b 1 x) = b := by
  have h := row_residual a x b 1 ha
  linarith

theorem distance_identity (a x z : Vec n) (b ω : ℝ)
    (ha : dot a a ≠ 0) (hz : dot a z = b) :
    sqdist (step a b ω x) z = sqdist x z -
      ω * (2 - ω) * (b - dot a x) ^ 2 / dot a a := by
  have hs (c : ℝ) : sqdist (fun i => x i + c * a i) z =
      sqdist x z + 2 * c * (dot a x - dot a z) + c ^ 2 * dot a a := by
    simp only [sqdist, dot, mul_sub, Finset.mul_sum, ← Finset.sum_sub_distrib,
      ← Finset.sum_add_distrib]
    apply Finset.sum_congr rfl
    intro i _
    ring
  unfold step
  rw [hs, hz]
  field_simp
  ring

theorem relaxed_descent (a x z : Vec n) (b ω : ℝ)
    (ha : a ≠ 0) (hz : dot a z = b) (hω : 0 ≤ ω) (hω' : ω ≤ 2) :
    sqdist (step a b ω x) z ≤ sqdist x z := by
  have hp : 0 < dot a a := lt_of_le_of_ne (dot_self_nonneg a)
    (Ne.symm (fun h => ha ((dot_self_eq_zero a).mp h)))
  rw [distance_identity a x z b ω (ne_of_gt hp) hz]
  have h : 0 ≤ ω * (2 - ω) * (b - dot a x) ^ 2 / dot a a := by positivity
  linarith

theorem strict_descent (a x z : Vec n) (b ω : ℝ)
    (ha : a ≠ 0) (hz : dot a z = b) (hω : 0 < ω) (hω' : ω < 2)
    (hr : dot a x ≠ b) : sqdist (step a b ω x) z < sqdist x z := by
  have hp : 0 < dot a a := lt_of_le_of_ne (dot_self_nonneg a)
    (Ne.symm (fun h => ha ((dot_self_eq_zero a).mp h)))
  rw [distance_identity a x z b ω (ne_of_gt hp) hz]
  have hr' : b - dot a x ≠ 0 := sub_ne_zero.mpr (Ne.symm hr)
  have h : 0 < ω * (2 - ω) * (b - dot a x) ^ 2 / dot a a := by positivity
  linarith

theorem fixed_iff (a x : Vec n) (b : ℝ) (ha : dot a a ≠ 0) :
    step a b 1 x = x ↔ dot a x = b := by
  constructor
  · intro h
    simpa only [h] using satisfies_row a x b ha
  · intro h
    funext i
    simp [step, h]

/-- Common fixed points of all row projections are precisely solutions. -/
theorem system_fixed_iff {m : ℕ} (A : Fin m → Vec n) (b : Fin m → ℝ)
    (x : Vec n) (hA : ∀ j, dot (A j) (A j) ≠ 0) :
    (∀ j, step (A j) (b j) 1 x = x) ↔ (∀ j, dot (A j) x = b j) := by
  exact forall_congr' (fun j => fixed_iff (A j) x (b j) (hA j))

#print axioms dot_self_nonneg
#print axioms dot_self_eq_zero
#print axioms row_residual
#print axioms satisfies_row
#print axioms distance_identity
#print axioms relaxed_descent
#print axioms strict_descent
#print axioms fixed_iff
#print axioms system_fixed_iff
end
end PaperToLean.Kaczmarz