/-
Copyright (c) 2026 TorchLean
Released under MIT license as described in the file LICENSE.
Authors: TorchLean Team
-/

module

public import Nemotron.Cache.Error
public import FloatLib.Floats.Formats.BinaryInterchange.Analysis.Error
public import FloatLib.Numerics.Quantization.Deterministic.Rational

/-!
# Reusing FloatLib rounding

The real-valued IEEE grid adapter below uses FloatLib's proved half-ULP bound. That grid has no
upper exponent limit; relating it to executable finite bits requires FloatLib's finiteness
refinement hypotheses. It is nearest-even rounding, whereas the released Ultra cache uses
stochastic rounding. No equivalence between those storage policies is asserted.
-/

@[expose] public section

namespace Nemotron.Cache

open Spec TorchLean KimiK3
open FloatLib.Floats
open FloatLib.Floats.Formats.BinaryInterchange

/-- Apply FloatLib's real-valued nearest-even grid independently to tensor coordinates. -/
noncomputable def roundTensor {shape : Shape} (fmt : FloatFormat)
    (x : Tensor ℝ shape) : Tensor ℝ shape :=
  Tensor.map (Model.roundAt fmt) x

/-- A supplied upper bound on half-ULP values discharges the tensor storage-error contract.
The premise includes subnormal values and does not assume a relative error bound near zero. -/
theorem roundTensor_within {shape : Shape} (fmt : FloatFormat)
    (x : Tensor ℝ shape) (δ : ℝ)
    (hδ : ∀ coordinate, Model.epsilonAt fmt (x coordinate) ≤ δ) :
    Within δ (roundTensor fmt x) x := by
  intro coordinate
  simpa [roundTensor, Tensor.map] using
    (Model.abs_roundAt_sub_le fmt (x coordinate)).trans (hδ coordinate)

namespace Counterexample

/-- The actual FloatLib integer-grid nearest-even quantizer, with rational interpretation. -/
def quantize (x : ℚ) : ℚ := (FloatLib.Numerics.roundRatEven x : ℚ)

/-- An exact rational recurrence isolates storage rounding from transition arithmetic. -/
def addStep (state input : ℚ) : ℚ × Unit := (state + input, ())

/-- A genuine uniform storage-error bound for this quantizer. -/
theorem quantize_error (x : ℚ) : |quantize x - x| ≤ (1 : ℚ) / 2 :=
  FloatLib.Numerics.roundRatEven_error_le_half x

/-- Four nearest-even writes can have less actual error than two writes, through cancellation.
This refutes a universal monotonic-error reading of "fewer sequential quantization steps".
It is a fixed-grid arithmetic example, not a counterexample to the report's measured results. -/
theorem fewer_writes_can_increase_error :
    |quantizedChunks addStep quantize 0
        [[-9 / 10], [-9 / 10], [-1 / 2], [-1 / 10]] - (-12 / 5)| <
      |quantizedChunks addStep quantize 0
        [[-9 / 10, -9 / 10], [-1 / 2, -1 / 10]] - (-12 / 5)| := by
  decide +kernel

/-- Repeated ties can suppress two half-unit writes, although one checkpoint preserves their sum. -/
theorem repeated_ties_lose_updates :
    quantizedChunks addStep quantize 0 [[1 / 2], [1 / 2]] = 0 ∧
      quantizedChunks addStep quantize 0 [[1 / 2, 1 / 2]] = 1 := by
  decide +kernel

end Counterexample

end Nemotron.Cache