import AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ReferenceCarryingKernel
import AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.JointReferenceGradientDescent
import Mathlib.Probability.Kernel.Composition.IntegralCompProd
import Mathlib.Tactic
/-!
# Actual reference queries under the recursive state law

Source: Chen, Chewi, Lu and Zhang, arXiv:2609.06906v1, Algorithm 3.3
and Section 6.3. The counted algorithm is exactly the reference-carrying
first gradient-square hit from the pre-noise current output. Completing
the source quadratic potential changes neither genuine gradients nor
any iterate, so it changes neither that first hit nor its count.

An active state charges N+1 visited gradient checks. An absorbed state
charges zero and does not execute reference descent. This is the stated
gradient-check reuse convention, not a verified evaluator or a discount
for gradients cached across different sampler calls.

The source current-call squared Wasserstein budget is required only for
active states and is in original coordinates. Genuine current curvature
and the real input moment imply the explicit logarithmic expected count;
neither a moment nor a count bound is assumed. The sampler achieving that
input precision is not constructed here. No unbounded cost is transferred
using total variation.

The joint integration measure is the composition-product of a state measure
and the current-output/independent-Gaussian product kernel. For a probability
state law this is a joint expectation; arbitrary SFinite state measures give
integral bounds. At an absorbed state the auxiliary product representation
has zero cost and does not assert that the stopped program draws randomness.
This is not itself the next-state marginal, nor proof that an arbitrary
input state law was generated by the full recursive execution. The nonnegative
extended-integral bound precedes the real expectation statement. Global
L1 is derived only when the displayed statewise upper bound is integrable.
No unconditional integrability, history invariant, stage sum, source
parameter schedule, stopping depth or terminal FORS claim follows.

Disclosed extensions include general beta, a coordinate-free positive
dimension, arbitrary retained states and any SFinite state measure.
The precision parameter epsilon may be any real-valued function: only
its square occurs. Measurability of epsilon implies measurability of the
displayed upper bound. Without it the extended integral uses Mathlib's
general nonnegative-integral convention, not an automatically measurable
random budget. The global real bound explicitly assumes its integrability.
The variances and base curvature remain strictly positive. Neither full
paper main result is asserted.
-/

open MeasureTheory ProbabilityTheory
open scoped ENNReal NNReal
noncomputable section
namespace AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ReferenceCarryingCost
variable {E : Type*} [NormedAddCommGroup E] [InnerProductSpace ℝ E] [CompleteSpace E]
omit [CompleteSpace E] in
private theorem regularized_hessian {U : E → ℝ} (hU : ContDiff ℝ 2 U)
    (r : ℝ) (u : E) :
    ContDiff ℝ 2 (fun x => U x+r/2*‖x-u‖^2) ∧
    ∀ x v, fderiv ℝ (fderiv ℝ (fun x => U x+r/2*‖x-u‖^2)) x v v =
      fderiv ℝ (fderiv ℝ U) x v v+r*‖v‖^2 := by
  let W := fun x => U x+r/2*‖x-u‖^2
  have hUd : Differentiable ℝ U := hU.differentiable (by norm_num)
  have hUdd : Differentiable ℝ (fderiv ℝ U) :=
    (hU.fderiv_right (m:=1) (by norm_num)).differentiable_one
  have hW : ContDiff ℝ 2 W := hU.add
    (contDiff_const.mul ((contDiff_id.sub contDiff_const).norm_sq (𝕜:=ℝ)))
  have hq (x : E) : HasFDerivAt (fun z => r/2*‖z-u‖^2)
      (r • innerSL ℝ (x-u)) x := by
    convert (((hasFDerivAt_id x).sub_const u).norm_sq).const_mul (r/2)
      using 1 <;> first | rfl | (ext v; simp; ring)
  have hfd (x : E) : fderiv ℝ W x = fderiv ℝ U x+r • innerSL ℝ (x-u) :=
    ((hUd x).hasFDerivAt.add (hq x)).fderiv
  let J : E →L[ℝ] (E →L[ℝ] ℝ) :=
    { toFun := fun v => innerSL ℝ v
      map_add' := by intros; ext; simp
      map_smul' := by intros; ext; simp
      cont := (innerSL ℝ (E:=E)).continuous }
  have hdd (x : E) : HasFDerivAt (fderiv ℝ W)
      (fderiv ℝ (fderiv ℝ U) x+r • J) x := by
    rw [show fderiv ℝ W = (fun z => fderiv ℝ U z+r • innerSL ℝ (z-u)) from funext hfd]
    convert (hUdd x).hasFDerivAt.add
      ((J.hasFDerivAt.comp x ((hasFDerivAt_id x).sub_const u)).const_smul r)
      using 1 <;> rfl
  refine ⟨hW,fun x v => ?_⟩
  rw [(hdd x).fderiv]
  change fderiv ℝ (fderiv ℝ U) x v v+r*inner ℝ v v = _
  rw [real_inner_self_eq_norm_sq]

private def firstIndex {X : Type*} (q : ℕ → X → ℝ) (b : X → ℝ) (p : X) : ℕ := by
  classical
  exact if h : ∃ n, q n p ≤ b p then Nat.find h else 0

private theorem statewise_cost [FiniteDimensional ℝ E] [MeasurableSpace E] [BorelSpace E]
    {V : E → ℝ} {α β : ℝ≥0} (hα : 0 < α) (hαβ : α ≤ β)
    (hV : ContDiff ℝ 2 V)
    (hH : ∀ x w : E, (α:ℝ)*‖w‖^2 ≤ fderiv ℝ (fderiv ℝ V) x w w ∧
      fderiv ℝ (fderiv ℝ V) x w w ≤ (β:ℝ)*‖w‖^2)
    (hd : 0 < (Module.finrank ℝ E : ℝ)) (b : ℝ≥0) (u : E)
    (η τ ε : ℝ) (hη : 0 < η) (hτ : 0 < τ)
    (ν : Measure E) [IsProbabilityMeasure ν]
    (hw : AutoSamplingTheory.TechnicalLemmas.Measure.WassersteinSpace.wassersteinDistance ν
      (AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing
        ((volume : Measure E).tilted (fun x => -(V x+(b:ℝ)/2*‖x-u‖^2)))
        (Real.sqrt (η/((β:ℝ)+b)))) ^ 2 ≤
      ENNReal.ofReal ((((β:ℝ)+b)/((α:ℝ)+b))*ε^2/((β:ℝ)+b))) :
    let d : ℝ := Module.finrank ℝ E
    let bs : ℝ := (β:ℝ)+b
    let a : ℝ := (α:ℝ)+b
    let k := bs/a
    let v := (η+τ)/bs
    let W := fun p : E × E => fun x => V x+(b:ℝ)/2*‖x-u‖^2+
      v⁻¹/2*‖x-(p.1+Real.sqrt (τ/bs) • p.2)‖^2
    let T := fun p x => x-(bs+v⁻¹)⁻¹ • gradient (W p) x
    let N := firstIndex (fun n p => ‖gradient (W p) ((T p)^[n] p.1)‖^2)
      (fun _ => (a+v⁻¹)*d)
    Integrable (fun p => (N p:ℝ)+1) (ν.prod (stdGaussian E)) ∧
      (∫ p, (N p:ℝ)+1 ∂ν.prod (stdGaussian E)) ≤
        2+((bs+v⁻¹)/(a+v⁻¹))*Real.log
          (1+4*(η+τ)*(1+η)+2*(η+τ)*k*ε^2/d+τ/(η+τ)) := by
  intro d bs a k v W T N
  have hb : 0 < bs := add_pos_of_pos_of_nonneg
    (lt_of_lt_of_le (show 0 < (α:ℝ) from hα) (show (α:ℝ) ≤ β from hαβ)) b.coe_nonneg
  have ha : 0 < a := add_pos_of_pos_of_nonneg (show 0 < (α:ℝ) from hα) b.coe_nonneg
  have hn : 0 ≤ k*ε^2/bs := by dsimp [k]; positivity
  let U := fun x => V x+(b:ℝ)/2*‖x-u‖^2
  have hreg := regularized_hessian hV (b:ℝ) u
  have hHs (x w : E) : ((α+b:ℝ≥0):ℝ)*‖w‖^2 ≤ fderiv ℝ (fderiv ℝ U) x w w ∧
      fderiv ℝ (fderiv ℝ U) x w w ≤ ((β+b:ℝ≥0):ℝ)*‖w‖^2 := by
    rw [hreg.2]
    simp only [NNReal.coe_add]
    constructor <;> nlinarith [(hH x w).1,(hH x w).2]
  have hA : η/bs+τ/bs=v := by dsimp [v]; ring
  have H := AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.JointReferenceGradientDescent.joint_reference_gradient_descent
    (α:=α+b) (β:=β+b) (add_pos_of_pos_of_nonneg hα (show 0 ≤ b from zero_le))
    (show α+b ≤ β+b from add_le_add hαβ le_rfl)
    hreg.1 hHs (div_pos hη hb).le (div_pos hτ hb) (Real.sqrt_nonneg (k*ε^2/bs)) hd ν
    (show AutoSamplingTheory.TechnicalLemmas.Measure.WassersteinSpace.wassersteinDistance ν
      (AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing
        ((volume : Measure E).tilted (fun x => -U x)) (Real.sqrt (η/bs))) ^ 2 ≤
        ENNReal.ofReal ((Real.sqrt (k*ε^2/bs))^2) by simpa only [Real.sq_sqrt hn] using hw)
  have hi : Integrable (fun p => (N p:ℝ)+1) (ν.prod (stdGaussian E)) := by
    have hi := H.2.2.2.1
    simp only [hA,NNReal.coe_add] at hi
    exact hi
  have he := H.2.2.2.2
  simp only [hA,NNReal.coe_add,Real.sq_sqrt hn] at he
  have heq : 1+v*(4*bs*d+4*bs^2*(η/bs)*d+2*bs^2*(k*ε^2/bs)+(τ/bs)*d/v^2)/d =
      1+4*(η+τ)*(1+η)+2*(η+τ)*k*ε^2/d+τ/(η+τ) := by
    have hd0 : d ≠ 0 := hd.ne'
    dsimp [v]
    field_simp [hb.ne',hd0,(add_pos hη hτ).ne']
    ring
  refine ⟨hi,?_⟩
  change (∫ p, (N p:ℝ)+1 ∂ν.prod (stdGaussian E)) ≤
    2+((bs+v⁻¹)/(a+v⁻¹))*Real.log _ at he
  rwa [heq] at he


private theorem integrate_actual_cost {S X : Type*} [MeasurableSpace S] [MeasurableSpace X]
    (K : Kernel S X) [IsMarkovKernel K] (C : S × X → ℝ) (B : S → ℝ)
    (hC : Measurable C) (hC0 : ∀ p, 0 ≤ C p)
    (hi : ∀ s, Integrable (fun x => C (s,x)) (K s))
    (hb : ∀ s, (∫ x, C (s,x) ∂K s) ≤ B s)
    (μ : Measure S) [SFinite μ] :
    (∫⁻ p, ENNReal.ofReal (C p) ∂μ.compProd K) ≤ ∫⁻ s, ENNReal.ofReal (B s) ∂μ ∧
      (Integrable B μ → Integrable C (μ.compProd K) ∧
        (∫ p, C p ∂μ.compProd K) ≤ ∫ s, B s ∂μ) := by
  have hE : Measurable (fun s => ∫ x, C (s,x) ∂K s) :=
    hC.stronglyMeasurable.integral_kernel_prod_right'.measurable
  have hE0 (s : S) : 0 ≤ ∫ x, C (s,x) ∂K s := integral_nonneg (fun x => hC0 (s,x))
  constructor
  · rw [Measure.lintegral_compProd hC.ennreal_ofReal]
    apply lintegral_mono
    intro s
    change (∫⁻ x, ENNReal.ofReal (C (s,x)) ∂K s) ≤ ENNReal.ofReal (B s)
    rw [← ofReal_integral_eq_lintegral_ofReal (hi s) (Filter.Eventually.of_forall (fun x => hC0 (s,x)))]
    exact ENNReal.ofReal_le_ofReal (hb s)
  · intro hB
    have hEi : Integrable (fun s => ∫ x, C (s,x) ∂K s) μ :=
      hB.mono' hE.aestronglyMeasurable (Filter.Eventually.of_forall fun s => by
        rw [Real.norm_eq_abs,abs_of_nonneg (hE0 s)]
        exact hb s)
    have heq : (fun s => ∫ x, ‖C (s,x)‖ ∂K s) = (fun s => ∫ x, C (s,x) ∂K s) := by
      funext s
      apply integral_congr_ae
      exact Filter.Eventually.of_forall fun x => Real.norm_of_nonneg (hC0 (s,x))
    have hCi : Integrable C (μ.compProd K) :=
      (Measure.integrable_compProd_iff hC.aestronglyMeasurable).2
        ⟨Filter.Eventually.of_forall hi,by simpa only [heq] using hEi⟩
    refine ⟨hCi,?_⟩
    rw [Measure.integral_compProd hCi]
    exact integral_mono hEi hB hb

local notation "RefState" E:arg => ℝ≥0 × E × E × ℕ × (ℕ → E)

theorem reference_carrying_cost [FiniteDimensional ℝ E] [MeasurableSpace E] [BorelSpace E]
    {V : E → ℝ} {α β : ℝ≥0} (hα : 0 < α) (hαβ : α ≤ β)
    (hV : ContDiff ℝ 2 V)
    (hH : ∀ x w : E, (α:ℝ)*‖w‖^2 ≤ fderiv ℝ (fderiv ℝ V) x w w ∧
      fderiv ℝ (fderiv ℝ V) x w w ≤ (β:ℝ)*‖w‖^2)
    (hd : 0 < (Module.finrank ℝ E : ℝ))
    (η τ ε : RefState E → ℝ) (hη : Measurable η) (hτ : Measurable τ)
    (hη0 : ∀ s, 0 < η s) (hτ0 : ∀ s, 0 < τ s)
    (M : Kernel (RefState E) E) [IsMarkovKernel M] (threshold : ℝ≥0)
    (hw : ∀ s : RefState E, s.1 < threshold →
      AutoSamplingTheory.TechnicalLemmas.Measure.WassersteinSpace.wassersteinDistance (M s)
        (AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing
          ((volume : Measure E).tilted (fun x => -(V x+(s.1:ℝ)/2*‖x-s.2.1‖^2)))
          (Real.sqrt (η s/((β:ℝ)+s.1)))) ^ 2 ≤
        ENNReal.ofReal ((((β:ℝ)+s.1)/((α:ℝ)+s.1))*(ε s)^2/((β:ℝ)+s.1))) :
    let d : ℝ := Module.finrank ℝ E
    let bs := fun s : RefState E => (β:ℝ)+s.1
    let a := fun s : RefState E => (α:ℝ)+s.1
    let v := fun s => (η s+τ s)/bs s
    let bp := fun s : RefState E => s.1+Real.toNNReal (v s)⁻¹
    let obs := fun p : RefState E × (E × E) => p.2.1+Real.sqrt (τ p.1/bs p.1) • p.2.2
    let center := fun p : RefState E × (E × E) => (bp p.1:ℝ)⁻¹ •
      ((p.1.1:ℝ) • p.1.2.1+(v p.1)⁻¹ • obs p)
    let F := fun p x => V x+(bp p.1:ℝ)/2*‖x-center p‖^2
    let T := fun p x => x-((β:ℝ)+bp p.1)⁻¹ • gradient (F p) x
    let q := fun n p => ‖gradient (F p) ((T p)^[n] p.2.1)‖^2
    let N := firstIndex q (fun p => ((α:ℝ)+bp p.1)*d)
    let C := fun p : RefState E × (E × E) => if threshold ≤ p.1.1 then 0 else (N p:ℝ)+1
    let B := fun s : RefState E => if threshold ≤ s.1 then 0 else
      2+(((β:ℝ)+bp s)/((α:ℝ)+bp s))*Real.log
        (1+4*(η s+τ s)*(1+η s)+2*(η s+τ s)*(bs s/a s)*(ε s)^2/d+τ s/(η s+τ s))
    let K := M ×ₖ Kernel.const (RefState E) (stdGaussian E)
    Measurable C ∧ (Measurable ε → Measurable B) ∧ (∀ p, 0 ≤ C p) ∧
      (∀ s, Integrable (fun z => C (s,z)) (K s) ∧ (∫ z, C (s,z) ∂K s) ≤ B s) ∧
      (∀ (μ : Measure (RefState E)) [SFinite μ],
        (∫⁻ p, ENNReal.ofReal (C p) ∂μ.compProd K) ≤ ∫⁻ s, ENNReal.ofReal (B s) ∂μ ∧
          (Integrable B μ → Integrable C (μ.compProd K) ∧
            (∫ p, C p ∂μ.compProd K) ≤ ∫ s, B s ∂μ)) := by
  classical
  intro d bs a v bp obs center F T q N C B K
  have R := AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ReferenceCarryingKernel.reference_carrying_kernel
    hα hαβ hV hH hd η τ hη hτ hη0 hτ0 M threshold
  have hbp (s : RefState E) : (bp s:ℝ)=(s.1:ℝ)+(v s)⁻¹ := (R.1 s).2
  have hNm : Measurable N := R.2.1
  have hCm : Measurable C := by
    apply Measurable.ite (measurableSet_le measurable_const (by fun_prop)) measurable_const
    exact (measurable_from_nat.comp hNm).add measurable_const
  have hC0 (p : RefState E × (E × E)) : 0 ≤ C p := by dsimp [C]; split_ifs <;> positivity
  have hBm (hε : Measurable ε) : Measurable B := by
    dsimp [B,bp,v,bs,a]
    apply Measurable.ite (measurableSet_le measurable_const (by fun_prop)) measurable_const
    fun_prop
  have hK : IsMarkovKernel K := by dsimp [K]; infer_instance
  have hKs (s : RefState E) : K s=(M s).prod (stdGaussian E) := by
    simp only [K,Kernel.prod_apply,Kernel.const_apply]
  have hp (s : RefState E) : Integrable (fun z => C (s,z)) (K s) ∧
      (∫ z, C (s,z) ∂K s) ≤ B s := by
    by_cases hs : threshold ≤ s.1
    · simp only [C,B,if_pos hs,integrable_const,integral_const,smul_zero,le_refl,and_self]
    · have Hs := statewise_cost hα hαβ hV hH hd s.1 s.2.1 (η s) (τ s) (ε s)
        (hη0 s) (hτ0 s) (M s) (hw s (lt_of_not_ge hs))
      let W := fun z : E × E => fun x => V x+(s.1:ℝ)/2*‖x-s.2.1‖^2+
        (v s)⁻¹/2*‖x-(z.1+Real.sqrt (τ s/bs s) • z.2)‖^2
      let Ts := fun z x => x-(bs s+(v s)⁻¹)⁻¹ • gradient (W z) x
      let qs := fun n z => ‖gradient (W z) ((Ts z)^[n] z.1)‖^2
      let Ns := firstIndex qs (fun _ => (a s+(v s)⁻¹)*d)
      have hc : bs s+(v s)⁻¹=(β:ℝ)+bp s := by rw [hbp]; dsimp [bs]; ring
      have ha : a s+(v s)⁻¹=(α:ℝ)+bp s := by rw [hbp]; dsimp [a]; ring
      have hq (n : ℕ) (z : E × E) : qs n z=q n (s,z) := by
        have hz := R.2.2.2.2.1 (s,z)
        have ht := hz.2.2.2.2.2.1 n
        have hg := hz.2.2.2.2.1
        change (fun x => x-((β:ℝ)+bp s)⁻¹ • gradient (W z) x)^[n] z.1 =
          (T (s,z))^[n] z.1 at ht
        dsimp only [qs,Ts]
        rw [hc,ht]
        exact congrArg (fun w : E => ‖w‖^2) (hg _)
      have hN (z : E × E) : Ns z=N (s,z) := by
        dsimp only [Ns,N,firstIndex]
        simp only [hq,ha]
      change Integrable (fun z => (Ns z:ℝ)+1) ((M s).prod (stdGaussian E)) ∧
        (∫ z, (Ns z:ℝ)+1 ∂(M s).prod (stdGaussian E)) ≤
          2+((bs s+(v s)⁻¹)/(a s+(v s)⁻¹))*Real.log _ at Hs
      simp_rw [hN] at Hs
      simpa only [C,B,if_neg hs,hKs,hc,ha] using Hs
  exact ⟨hCm,hBm,hC0,hp,fun μ => integrate_actual_cost K C B hCm hC0
    (fun s => (hp s).1) (fun s => (hp s).2) μ⟩

end AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ReferenceCarryingCost