Tests.SmoothedPicardHMCActualIncrement

import AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ActualIncrement
import Mathlib.Analysis.SpecialFunctions.Trigonometric.Deriv

noncomputable section
set_option backward.isDefEq.respectTransparency false
open Set MeasureTheory ProbabilityTheory InnerProductSpace
open scoped BigOperators RealInnerProductSpace
open AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ActualIncrement

-- Nonconstant positive Hessian, not a supplied gradient/Hessian certificate.
private def potential (x : ℝ) : ℝ := (3/8)*x^2 + (1/8)*Real.sin x
private theorem potential_curvature :
    ∀ z v : ℝ, (1/(2*(1:ℝ)))*‖v‖^2 ≤ (fderiv ℝ (fderiv ℝ potential) z v) v ∧
      (fderiv ℝ (fderiv ℝ potential) z v) v ≤ ‖v‖^2 := by
  have hder (z : ℝ) : HasDerivAt potential ((3/4)*z+(1/8)*Real.cos z) z := by
    convert (((hasDerivAt_id z).pow 2).const_mul (3/8)).add
      ((Real.hasDerivAt_sin z).const_mul (1/8)) using 1 <;> first | rfl | (simp [potential, Pi.add_apply, mul_comm]; ring)
  have hD : fderiv ℝ potential = fun z =>
      ((3/4)*z+(1/8)*Real.cos z) • ContinuousLinearMap.id ℝ ℝ := by
    funext z
    ext
    rw [fderiv_eq_smul_deriv, (hder z).deriv]
    simp
  have hDD (z v : ℝ) : (fderiv ℝ (fderiv ℝ potential) z v) v =
      ((3/4)-(1/8)*Real.sin z)*v^2 := by
    have hd : HasDerivAt (fun w : ℝ => (3/4)*w+(1/8)*Real.cos w)
        ((3/4)-(1/8)*Real.sin z) z := by
      convert ((hasDerivAt_id z).const_mul (3/4)).add
        ((Real.hasDerivAt_cos z).const_mul (1/8)) using 1 <;> (try dsimp only [id]) <;> first | rfl | ring
    rw [hD, fderiv_eq_smul_deriv,
      (hd.smul_const (ContinuousLinearMap.id ℝ ℝ)).deriv]
    simp
    ring
  intro z v
  rw [hDD]
  simp only [Real.norm_eq_abs, sq_abs]
  constructor <;> nlinarith [Real.sin_le_one z, Real.neg_one_le_sin z, sq_nonneg v]

open AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.PhaseMetric

-- Genuine nonconstant Hessian and actual full phase, for every shared noise/input.
example {h : ℝ} {J : ℕ} (hJ : 2 ≤ J) (hh : 0 < h) :
    let t := fun i : Fin J => h/2*(1-Real.cos ((i : ℝ)/(J-1 : ℝ)*Real.pi))
    let ell := fun j : Fin J => Lagrange.basis Finset.univ t j
    let Λ := sSup ((fun s : ℝ => ∑ j : Fin J, |(ell j).eval s|) '' Icc 0 h)
    let ω := fun i j : Fin J => ∫ s in 0..t i, (t i-s)*(ell j).eval s
    let b := fun j : Fin J => ∫ s in 0..h, (ell j).eval s
    let c := fun j : Fin J => ∫ s in 0..h, (h-s)*(ell j).eval s
    h * Λ ≤ 1/(65536*(1:ℝ)) →
    let a := Real.exp (-h / 2)
    let sigma := Real.sqrt (1 - Real.exp (-h))
    let P0 := fun (z ζ : ℝ × ℝ) => a • z.2 + sigma • ζ.1
    let Y0 := fun (z ζ : ℝ × ℝ) j => z.1 + t j • P0 z ζ
    let Y1 := fun (z ζ : ℝ × ℝ) i => Y0 z ζ i - ∑ j, ω i j • gradient potential (Y0 z ζ j)
    let Φ := fun w : (ℝ × ℝ) × (ℝ × ℝ) =>
      (w.1.1 + h • P0 w.1 w.2 - ∑ j, c j • gradient potential (Y1 w.1 w.2 j),
        a • (P0 w.1 w.2 - ∑ j, b j • gradient potential (Y1 w.1 w.2 j)) + sigma • w.2.2)
    ∀ z z' ζ : ℝ × ℝ,
      Real.sqrt (phaseQuadraticForm (1:ℝ) (((Φ (z,ζ)).1-(Φ (z',ζ)).1)-(z.1-z'.1),
        ((Φ (z,ζ)).2-(Φ (z',ζ)).2)-(z.2-z'.2))) ≤
      12*h*Real.sqrt (phaseQuadraticForm (1:ℝ) (z.1-z'.1,z.2-z'.2)) := by
  have hf : ContDiff ℝ 2 potential := by unfold potential; fun_prop
  exact source_synchronous_increment (by norm_num : (1:ℝ) ≤ 1) hf potential_curvature hJ hh

#check @source_synchronous_increment
#print axioms source_synchronous_increment
#print axioms AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.PhaseMetric.ordinary_energy_le_phaseQuadraticForm
#print axioms AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.PhaseMetric.phaseQuadraticForm_le_ordinary_energy