import Mathlib.LinearAlgebra.BilinearForm.Properties
import Mathlib.Data.Real.Basic
import Mathlib.Tactic.FieldSimp
import Mathlib.Tactic.Ring
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.Positivity

/-! Frisch and Waugh (1933): one nuisance regressor and one target regressor.
An exact loss decomposition proves residualization gives the full-regression
coefficient. Abstract real bilinear forms include ordinary finite data vectors.
No intercept is implicit; include or remove one explicitly when using the result. -/
namespace PaperToLean.FrischWaugh
noncomputable section
open LinearMap (BilinForm)
variable {V : Type*} [AddCommGroup V] [Module ℝ V]

def residual (B : BilinForm ℝ V) (z v : V) : V := v - (B z v / B z z) • z
def loss (B : BilinForm ℝ V) (y x z : V) (a b : ℝ) : ℝ :=
  B (y - a • x - b • z) (y - a • x - b • z)
def nuisance (B : BilinForm ℝ V) (y x z : V) (a : ℝ) : ℝ :=
  (B z y - a * B z x) / B z z
def coefficient (B : BilinForm ℝ V) (y x z : V) : ℝ :=
  B (residual B z x) (residual B z y) / B (residual B z x) (residual B z x)

theorem residual_orthogonal (B : BilinForm ℝ V) (z v : V) (hz : B z z ≠ 0) :
    B z (residual B z v) = 0 := by
  simp only [residual, BilinForm.sub_right, BilinForm.smul_right]
  field_simp
  ring

theorem loss_split (B : BilinForm ℝ V) (hB : B.IsSymm) (y x z : V)
    (hz : B z z ≠ 0) (a b : ℝ) :
    loss B y x z a b =
      B (residual B z y - a • residual B z x) (residual B z y - a • residual B z x) +
      B z z * (b - nuisance B y x z a) ^ 2 := by
  simp only [loss, residual, nuisance, BilinForm.sub_left, BilinForm.sub_right,
    BilinForm.smul_left, BilinForm.smul_right]
  rw [hB.eq y z, hB.eq x z, hB.eq y x]
  field_simp
  ring

theorem best_nuisance (B : BilinForm ℝ V) (hB : B.IsSymm) (y x z : V)
    (hz : 0 < B z z) (a b : ℝ) :
    loss B y x z a (nuisance B y x z a) ≤ loss B y x z a b := by
  rw [loss_split B hB y x z (ne_of_gt hz), loss_split B hB y x z (ne_of_gt hz)]
  have h : 0 ≤ B z z * (b - nuisance B y x z a) ^ 2 := by positivity
  nlinarith

theorem one_variable_square (B : BilinForm ℝ V) (hB : B.IsSymm) (u v : V)
    (hu : B u u ≠ 0) (a : ℝ) :
    B (v - a • u) (v - a • u) =
      B v v - (B u v) ^ 2 / B u u + B u u * (a - B u v / B u u) ^ 2 := by
  simp only [BilinForm.sub_left, BilinForm.sub_right, BilinForm.smul_left,
    BilinForm.smul_right]
  rw [hB.eq v u]
  field_simp
  ring

/-- Full loss = optimum loss + two explicit nonnegative penalties. -/
theorem full_loss_gap (B : BilinForm ℝ V) (hB : B.IsSymm) (y x z : V)
    (hz : B z z ≠ 0) (hx : B (residual B z x) (residual B z x) ≠ 0) (a b : ℝ) :
    loss B y x z a b =
      loss B y x z (coefficient B y x z) (nuisance B y x z (coefficient B y x z)) +
      B (residual B z x) (residual B z x) * (a - coefficient B y x z) ^ 2 +
      B z z * (b - nuisance B y x z a) ^ 2 := by
  rw [loss_split B hB y x z hz, loss_split B hB y x z hz]
  rw [one_variable_square B hB _ _ hx, one_variable_square B hB _ _ hx]
  simp only [coefficient, sub_self, zero_pow (by decide : 2 ≠ 0), mul_zero, add_zero]

theorem unique_coefficients (B : BilinForm ℝ V) (hB : B.IsSymm) (y x z : V)
    (hz : 0 < B z z) (hx : 0 < B (residual B z x) (residual B z x)) (a b : ℝ) :
    loss B y x z a b ≤
      loss B y x z (coefficient B y x z) (nuisance B y x z (coefficient B y x z)) ↔
      a = coefficient B y x z ∧ b = nuisance B y x z (coefficient B y x z) := by
  constructor
  · intro h
    have hg := full_loss_gap B hB y x z (ne_of_gt hz) (ne_of_gt hx) a b
    have hp : 0 ≤ B z z * (b - nuisance B y x z a) ^ 2 := by positivity
    have ha : (a - coefficient B y x z) ^ 2 = 0 := by
      nlinarith [sq_nonneg (a - coefficient B y x z)]
    have hae := sub_eq_zero.mp (sq_eq_zero_iff.mp ha)
    have hb : (b - nuisance B y x z a) ^ 2 = 0 := by
      nlinarith [sq_nonneg (b - nuisance B y x z a)]
    exact ⟨hae, by simpa [hae] using sub_eq_zero.mp (sq_eq_zero_iff.mp hb)⟩
  · rintro ⟨rfl, rfl⟩
    exact le_rfl

#print axioms residual_orthogonal
#print axioms loss_split
#print axioms best_nuisance
#print axioms one_variable_square
#print axioms full_loss_gap
#print axioms unique_coefficients
end
end PaperToLean.FrischWaugh