import FloatLib.Numerics.Enclosure.Elementary.Proof
import TorchLean.NN.Softmax.Correctness
import TorchLean.Verify.RuntimeApprox.Roundoff.CrossEntropy
import TorchLeanExamples.DeepDives.Floats.RoundingBugs

/-!
## Definitions and the exact real-valued gradient

Checked with Lean 4.35.0 and TorchLean revision bb95db148ab63e3beb3e9c43813ad782bdddad67.
Run `lake env lean /path/to/SoftmaxBlog.lean` from that TorchLean checkout.
-/

namespace TorchLean.Examples.BlogSoftmax

open Tensor
open FloatLib.Floats FloatLib.Numerics
open FloatLib.Floats.Formats.BinaryInterchange

section Formulas

variable {α : Type} [Storage α] [Context α]

/-- Logits `[m, -m, -m]`. The target class `0` leads the other two by `2 * m`. -/
def logits (m : α) : Tensor α [3] := Tensor.ofFn fun i => if i = 0 then m else -m

/-- The target entry of the standard log-softmax at `logits m`. -/
def standardLogProbability (m : α) : α :=
  (StandardSoftmax.logSoftmax 0 (logits m)).getScalar 0

/-- The standard cross-entropy of `logits m` against the target `0`. -/
def standardLoss (m : α) : α :=
  StandardSoftmax.crossEntropyLogits 0 (logits m) SoftmaxCrossEntropy.target

/-- The target entry of the standard cross-entropy gradient at `logits m`. -/
def standardGradient (m : α) : α :=
  (StandardSoftmax.crossEntropyLogitsDeriv 0 (logits m)
    SoftmaxCrossEntropy.target).getScalar 0

/-- The target entry of `Tensor.crossEntropyBackward .logits` at `logits m`. -/
def gradient (m : α) : α :=
  (Tensor.crossEntropyBackward .logits (logits m) SoftmaxCrossEntropy.target ⟨0,
      by exact Nat.zero_lt_succ _⟩).getScalar 0

end Formulas

/-- The softmax probability `2 exp (-gap) / (1 + 2 exp (-gap))` of the two classes of
`logits (gap / 2)` other than the target. -/
noncomputable def tailMass (gap : ℝ) : ℝ :=
  2 * Real.exp (-gap) / (1 + 2 * Real.exp (-gap))

/-- Over the reals, the target gradient at `logits m` is `-tailMass (2 * m)`. -/
theorem gradient_real (m : ℝ) : gradient m = -tailMass (2 * m) := by
  have connected := congrArg (fun value : Vec 3 => value 0)
    (Tensor.getScalarE_logSoftmaxBackward (logits m)
      (Tensor.map (fun x => -x) SoftmaxCrossEntropy.target))
  have coordinate : gradient m = -1 + Real.exp m / (Real.exp m + 2 * Real.exp (-m)) := by
    have hdenom : axisMeanDenom [3] ⟨0, Shape.AxisInBounds.proof⟩ = 1 := rfl
    simp only [gradient, Tensor.crossEntropyBackward, CrossEntropy.logProbs_real, reduceIte, hdenom]
    simp [Vec.logSoftmaxVJP, Vec.softmax, Vec.sumExp,
      Fin.sum_univ_three, logits, SoftmaxCrossEntropy.target] at connected ⊢
    rw [connected]
    ring
  have hexp : Real.exp (-(2 * m)) = Real.exp (-m) * Real.exp (-m) := by
    rw [← Real.exp_add]
    ring_nf
  have hinverse : Real.exp m * Real.exp (-m) = 1 := by
    rw [← Real.exp_add, add_neg_cancel, Real.exp_zero]
  rw [coordinate, tailMass, hexp]
  have hpositive : 0 < Real.exp (-m) := Real.exp_pos _
  field_simp
  nlinarith [hinverse]

end TorchLean.Examples.BlogSoftmax

/-! ## The binary32 counterexample and the replacement's error bound -/

namespace TorchLean.Examples.BlogSoftmax

open Tensor
open FloatLib.Floats FloatLib.Numerics
open FloatLib.Floats.Formats.BinaryInterchange

/-- Rational bounds `lo` and `hi` on the endpoints of FloatLib's enclosure of `exp (-gap)` bound
`tailMass gap` by `2 lo / (1 + 2 lo)` and `2 hi / (1 + 2 hi)`, so by any `a` below the
first and any `b` above the second. -/
private theorem tailMass_mem_Icc (gap degree : ℕ) (lo hi : ℚ) {a b : ℝ}
    (hlo₀ : 0 ≤ lo) (hlo : lo ≤ (Enclosure.exp (-gap) degree).lo)
    (hhi : (Enclosure.exp (-gap) degree).hi ≤ hi)
    (ha : a ≤ ((2 * lo / (1 + 2 * lo) : ℚ) : ℝ)) (hb : ((2 * hi / (1 + 2 * hi) : ℚ) : ℝ) ≤ b) :
    tailMass gap ∈ Set.Icc a b := by
  obtain ⟨lower, upper⟩ := Enclosure.contains_exp (-gap) degree
  have increasing {s t : ℝ} (hs : 0 ≤ s) (hst : s ≤ t) :
      2 * s / (1 + 2 * s) ≤ 2 * t / (1 + 2 * t) := by
    rw [div_le_div_iff₀ (by positivity) (by linarith)]
    nlinarith
  push_cast at lower upper ha hb ⊢
  exact ⟨ha.trans (increasing (by exact_mod_cast hlo₀) ((Rat.cast_le.mpr hlo).trans lower)),
    (increasing (Real.exp_pos _).le (upper.trans (Rat.cast_le.mpr hhi))).trans hb⟩

section Bounds

variable {format : FloatFormat} {plan : Configured.StoragePlan format} {code : Type}
    [ExecFloat.ModelCodec plan (Model format) code]

local notation "Value" =>
  ExecFloat.Binary format.expWidth format.fracWidth format.encoding format.exponentBias
    format.expWidth_ge_two format.fracWidth_pos format.exponentBias_pos
    format.exponentBias_le_maxFinite plan code

/-- A value that decodes to zero is not within 32 units of roundoff and underflow of `-c`, for
`c ∈ [lo, hi]` with `0 < lo`, when `32 (u hi + 2 ^ e)` is below `lo` for a power `2 ^ e` at least
the underflow allowance. -/
private theorem not_mixedError_of_eq_zero {x : Value} {c lo hi : ℝ} (e : ℤ)
    (hx : ExecFloat.Binary.toRat? x = some 0) (hc : c ∈ Set.Icc lo hi) (hlo₀ : 0 < lo)
    (he : FloatFormat.minSubnormalExponent format - 1 ≤ e)
    (hallowance : 32 * (Model.unitRoundoffAt format * hi + 2 ^ e) < lo) :
    ¬ MixedError (ExecFloat.Binary.toReal x) (-c) (32 * Model.unitRoundoffAt format)
      (32 * Model.underflowErrorAt format) := by
  rw [ExecFloat.Binary.toReal_of_toRat?_eq_some hx, Rat.cast_zero]
  intro bound
  obtain ⟨hlo, hhi⟩ := hc
  have hu := Model.unitRoundoffAt_pos format
  have hη := Model.underflowErrorAt_le_zpow format he
  have hscaled := mul_le_mul_of_nonneg_left hhi hu.le
  have hpositive : 0 < c := hlo₀.trans_le hlo
  have h := bound.bound
  rw [zero_sub, neg_neg, abs_neg, abs_of_pos hpositive] at h
  linarith

/-- A value that decodes to a negative `q` is finite, negative and within 32 units of roundoff
and underflow of `-c`, for `c ∈ [lo, hi]`, when `-32 u lo ≤ q + lo` and `q + hi ≤ 32 u lo`. -/
private theorem mixedError_of_eq_some {x : Value} {q : ℚ} {c lo hi : ℝ}
    (hx : ExecFloat.Binary.toRat? x = some q) (hq : q < 0) (hc : c ∈ Set.Icc lo hi)
    (hlower : -(32 * Model.unitRoundoffAt format * lo) ≤ q + lo)
    (hupper : q + hi ≤ 32 * Model.unitRoundoffAt format * lo) :
    ExecFloat.Binary.isFinite x = true ∧
      MixedError (ExecFloat.Binary.toReal x) (-c) (32 * Model.unitRoundoffAt format)
        (32 * Model.underflowErrorAt format) ∧
      ExecFloat.Binary.toReal x < 0 := by
  obtain ⟨hlo, hhi⟩ := hc
  have hu := Model.unitRoundoffAt_pos format
  have hscaled : 32 * Model.unitRoundoffAt format * lo ≤ 32 * Model.unitRoundoffAt format * |-c| :=
    mul_le_mul_of_nonneg_left (hlo.trans ((le_abs_self c).trans (abs_neg c).ge)) (by positivity)
  rw [ExecFloat.Binary.toReal_of_toRat?_eq_some hx]
  refine ⟨Model.Algebraic.isFinite_of_toRat?_eq_some hx,
    (MixedError.of_abs_sub_le_mul (by positivity) ?_).mono le_rfl
      (mul_nonneg (by norm_num) (Model.underflowErrorAt_pos format).le),
    by exact_mod_cast hq⟩
  rw [sub_neg_eq_add, abs_le]
  constructor <;> linarith

end Bounds

/-! ## Binary32

The paper's logits `[9, -9, -9]`. With a gap of `18`, `2 exp (-18)` is below the binary32 unit
roundoff `2^-24`. -/

/-- FloatLib's enclosure of `exp (-18)` gives `tailMass 18` to eight digits. -/
private theorem binary32_tailMass_mem_Icc :
    tailMass 18 ∈ Set.Icc (30459958 / 10 ^ 15) (30459959 / 10 ^ 15) :=
  tailMass_mem_Icc 18 30 (152299797 / 10 ^ 16) (152299798 / 10 ^ 16)
    (by norm_num) (by decide +kernel) (by decide +kernel) (by norm_num) (by norm_num)

/-- The standard formulas evaluated with the binary32 primitives return zero for the target
log-probability, the loss and the target gradient. -/
theorem binary32_standard_toRat? :
    ExecFloat.Binary.toRat? (standardLogProbability (9 : Binary32)) = some 0 ∧
      ExecFloat.Binary.toRat? (standardLoss (9 : Binary32)) = some 0 ∧
      ExecFloat.Binary.toRat? (standardGradient (9 : Binary32)) = some 0 := by
  decide +kernel

/-- The exact probability of the non-target classes lies between `3e-8` and `3.1e-8`. -/
theorem binary32_tailMass_mem_Ioo :
    tailMass 18 ∈ Set.Ioo (3 / 10 ^ 8) (31 / 10 ^ 9) := by
  obtain ⟨lower, upper⟩ := binary32_tailMass_mem_Icc
  exact ⟨lt_of_lt_of_le (by norm_num) lower, lt_of_le_of_lt upper (by norm_num)⟩

/-- An absolute tolerance of `1e-7` accepts the standard target gradient. -/
theorem binary32_abs_sub_lt :
    |ExecFloat.Binary.toReal (standardGradient (9 : Binary32)) - -tailMass 18| <
      1 / 10 ^ 7 := by
  obtain ⟨lower, upper⟩ := binary32_tailMass_mem_Ioo
  rw [ExecFloat.Binary.toReal_of_toRat?_eq_some binary32_standard_toRat?.2.2, Rat.cast_zero,
    zero_sub, neg_neg, abs_of_pos (by linarith)]
  linarith

/-- The mixed bound with 32 units of binary32 roundoff and underflow rejects the standard target
gradient. -/
theorem binary32_not_mixedError :
    ¬ MixedError (ExecFloat.Binary.toReal (standardGradient (9 : Binary32))) (-tailMass 18)
      (32 * Model.unitRoundoffAt FloatFormat.binary32)
      (32 * Model.underflowErrorAt FloatFormat.binary32) :=
  not_mixedError_of_eq_zero (-64) binary32_standard_toRat?.2.2
    binary32_tailMass_mem_Icc (by norm_num) (by decide)
    (by norm_num [Model.unitRoundoffAt_eq, FloatFormat.binary32])

/-- The target gradient `gradient 9` is finite, negative and within the mixed bound with 32
units of binary32 roundoff and underflow. -/
theorem binary32_mixedError :
    ExecFloat.Binary.isFinite (gradient (9 : Binary32)) = true ∧
      MixedError (ExecFloat.Binary.toReal (gradient (9 : Binary32))) (-tailMass 18)
        (32 * Model.unitRoundoffAt FloatFormat.binary32)
        (32 * Model.underflowErrorAt FloatFormat.binary32) ∧
      ExecFloat.Binary.toReal (gradient (9 : Binary32)) < 0 :=
  mixedError_of_eq_some (q := -2143429 / 2 ^ 46) (by decide +kernel) (by norm_num)
    binary32_tailMass_mem_Icc (by norm_num [Model.unitRoundoffAt_eq, FloatFormat.binary32])
    (by norm_num [Model.unitRoundoffAt_eq, FloatFormat.binary32])

end TorchLean.Examples.BlogSoftmax

/-! ## Generating an error bound for variable inputs -/

namespace TorchLean.Examples.BlogSoftmaxRoundoff

open Tensor
open FloatLib.Floats FloatLib.Numerics
open FloatLib.Floats.Formats.BinaryInterchange

variable {format : FloatFormat} {plan : Configured.StoragePlan format} {code : Type}
  [ExecFloat.ModelCodec plan (Model format) code]

local notation "Value" => ExecFloat (Configured.Family format code plan)

attribute [local roundoff_rule] TorchLean.Roundoff.CrossEntropy.log_probs_error

/-- The cross-entropy gradient's schedule on decoded logits `x`, with backward pivot `p'` and the
sum `S` of the exponentials `exp (x j)` folded in coordinate order, has the exact gradient
`softmax x i - [i = label]` as its reference: `exp (x j - log S)` is `softmax x j`, and at the
backward pivot the VJP's formula for the pivot entry has the same value. -/
theorem softmax_sub_oneHot_eq {n : Nat} (x : Fin (n + 1) → ℝ) (p' label i : Fin (n + 1)) :
    let S := (Tensor.Rep.coordinates [n + 1]).foldl (fun sum coordinate =>
      sum + Real.exp (x coordinate.1)) 0
    let q (j : Fin (n + 1)) : ℝ := Real.exp (x j - Real.log S)
    (if i = p' then
        -(q p' - 1) * -(if i = label then 1 else 0) - q p' * if label = p' then 0 else -1
      else q i - if i = label then 1 else 0) =
      Real.exp (x i) / ∑ j, Real.exp (x j) - if i = label then 1 else 0 := by
  intro S q
  have hS : S = ∑ j, Real.exp (x j) := Tensor.Rep.foldl_coordinates_add fun j => Real.exp (x j)
  have hq (j : Fin (n + 1)) : q j = Real.exp (x j) / ∑ k, Real.exp (x k) := by
    change Real.exp (x j - Real.log S) = _
    rw [Real.exp_sub, hS,
      Real.exp_log (Finset.sum_pos (fun _ _ => Real.exp_pos _) Finset.univ_nonempty)]
  simp only [hq]
  split_ifs <;> simp_all

/-- `softmax_sub_oneHot_eq` on the decoded logits of a vector: the reference `roundoff` infers for
the cross-entropy gradient is the exact gradient. -/
theorem gradient_reference [Storage Value] {n : Nat} (t : Tensor Value [n + 1])
    (label i : Fin (n + 1)) :
    let p' : Fin (n + 1) := Softmax.pivot (CrossEntropy.vectorLogProbs t)
    let x (j : Fin (n + 1)) : ℝ :=
      ExecFloat.Binary.toReal (Tensor.Rep.get t (j, PUnit.unit))
    let q (j : Fin (n + 1)) : ℝ := Real.exp (x j - Real.log
      ((Tensor.Rep.coordinates [n + 1]).foldl (fun sum coordinate =>
        sum + Real.exp (x coordinate.1)) 0))
    (if i = p' then
        -(q p' - 1) * -(if i = label then 1 else 0) - q p' * if label = p' then 0 else -1
      else q i - if i = label then 1 else 0) =
      Real.exp (ExecFloat.Binary.toReal (t.getScalar i)) /
          ∑ j, Real.exp (ExecFloat.Binary.toReal (t.getScalar j)) -
        if i = label then 1 else 0 := by
  intro p' x q
  simp only [Tensor.getScalar_eq_apply]
  exact softmax_sub_oneHot_eq x p' label i

roundoff mixed theorem gradient_error using gradient_reference :=
  fun [Storage Value] {n : Nat} (t : Tensor Value [n + 1]) (label : Fin (n + 1))
    (_hformat : format.isIEEE = true)
    (_hfinite : ∀ i, ExecFloat.Binary.isFinite (t.getScalar i) = true) (i : Fin (n + 1)) =>
    ExecFloat.Binary.toReal ((Tensor.crossEntropyBackward .logits t
      (Tensor.ofFn fun j => if j = label then 1 else 0)
      ⟨0, by exact Nat.zero_lt_succ _⟩).getScalar i)

end TorchLean.Examples.BlogSoftmaxRoundoff