-- ============================================================================
-- SPDX-License-Identifier: GPL-3.0-or-later
-- Copyright (C) 2026 Alexander Allan (MDMAchine) -- A&E Concepts
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-- GNU General Public License for more details: https://www.gnu.org/licenses/
-- ============================================================================

-- MD PingPong Simple v1.1 — Ancestral Euler + Momentum + Look-Back Smoother
-- MDMAchine | A&E Concepts © 2026
--
-- Port of MD_PingPong_Samplers.py (single-branch path) to HOT-Step-CPP Lua solver.
--
-- WHAT THIS DOES:
--   Standard ODE solvers (Euler, Heun, STORM) advance the latent deterministically.
--   PingPong injects ancestral (stochastic) noise at each step — the "ping" is
--   the clean denoising step, the "pong" is the noise re-injection that keeps
--   the trajectory alive and stochastic.
--
--   CORE STEP MATH:
--     dt     = t_prev - t_curr               (negative — stepping down)
--     x_euler = xt + dt * vt                 (standard Euler, matches STORM/OmniRelational)
--     x_next  = x_euler + noise * |dt| * ancestral_strength
--
--   MOMENTUM:
--   Latent velocity (x - x_prev) carried forward at each step. Maintains
--   flow continuity across the ODE trajectory — reduces erratic jumps between
--   steps, especially at low step counts.
--
--   NOISE COHERENCE:
--   Blends fresh Gaussian noise with the previous step's noise at ratio
--   noise_coherence. 0=fully fresh, 1=fully carried. Useful for temporal
--   smoothness in audio; keep low (0-0.2) to avoid spectral smearing.
--
--   LOOK-BACK SNR SMOOTHER:
--   λ(σ) = lambda_base * (σ/σ_max)^snr_power — heavy at high sigma, fades to
--   zero at low sigma. Suppresses ODE manifold shearing and harmonic hum.
--   Reference: arXiv:2602.09449
--
--   RMS SERVO:
--   Downward-only energy ceiling that follows a smooth curve from rms_max
--   (high sigma) to rms_min (low sigma). Domain-tunable: image latents
--   typically sit around 0.75-0.97; audio latents around 0.3-0.7.
--   Servo is DOWNWARD ONLY — never boosts energy, only clamps excess.
--
-- SOLVER API NOTE:
--   HOT-Step-CPP passes velocity vt where dt = t_prev - t_curr is NEGATIVE
--   (stepping from high sigma to low sigma). Euler update is:
--     x_next = xt + dt * vt   (same convention as STORM and OmniRelational)
--   Ancestral noise is added as: noise * |dt| * strength
--   Do NOT use xt - t_curr * vt for denoised — sign convention mismatch.
--
-- PARAMS:
--   ancestral_strength — noise injection scale. 1.0=standard ancestral, 0=pure ODE
--   noise_coherence    — step-to-step noise carry. 0=fresh, 0.2=subtle temporal link
--   momentum_strength  — latent velocity carry-over. 0.15=subtle, 0.3=strong
--   look_back_enabled  — SNR smoother toggle
--   look_back_lambda   — max smoothing weight (0.55=25-step, 0.35=35-step)
--   look_back_snr_power — falloff exponent (1.3=25-step, 1.5=35-step)
--   rms_servo          — energy ceiling toggle
--   rms_target_min     — servo floor at low sigma (audio: ~0.3, image: ~0.75)
--   rms_target_max     — servo ceiling at high sigma (audio: ~0.7, image: ~0.97)
--   rms_servo_gain     — correction aggressiveness (0.6=default, 1.0=hard snap)
--   seed               — RNG seed
-- ============================================================================

solver = {
    name        = "md_pingpong_simple",
    display     = "MD PingPong Simple (Ancestral)",
    description = "Ancestral Euler with momentum, noise coherence, look-back SNR smoother, and domain-tunable RMS servo. Single-branch stochastic sampler. Port of MD_PingPong_Samplers v3.5.",
    nfe         = 1,
    order       = 1,
    needs_model = false,
    stateful    = true,
    stochastic  = true,
    params      = {
        {
            key     = "ancestral_strength",
            type    = "slider",
            label   = "Ancestral Strength",
            default = 0.2,
            min     = 0.0,
            max     = 1.5,
            step    = 0.05,
            hint    = "Noise injection strength. 1.0=standard ancestral. 0=pure ODE (no noise).",
        },
        {
            key     = "noise_coherence",
            type    = "slider",
            label   = "Noise Coherence",
            default = 0.0,
            min     = 0.0,
            max     = 1.0,
            step    = 0.05,
            hint    = "Step-to-step noise correlation. 0=fresh noise each step. 0.2=subtle temporal link. Keep low for audio to avoid smearing.",
        },
        {
            key     = "momentum_strength",
            type    = "slider",
            label   = "Momentum",
            default = 0.1,
            min     = 0.0,
            max     = 0.5,
            step    = 0.01,
            hint    = "Latent velocity carry-over. 0.1=subtle flow continuity. 0.3=strong.",
        },
        {
            key     = "look_back_enabled",
            type    = "toggle",
            label   = "Look-Back Smoother",
            default = true,
            hint    = "SNR-adaptive latent EMA. Suppresses ODE manifold shearing and harmonic hum. arXiv:2602.09449.",
        },
        {
            key     = "look_back_lambda",
            type    = "slider",
            label   = "Look-Back Lambda",
            default = 0.55,
            min     = 0.1,
            max     = 1.0,
            step    = 0.05,
            hint    = "Max smoothing weight. Active when Look-Back Smoother is on. 0.55=25-step, 0.35=35-step.",
        },
        {
            key     = "look_back_snr_power",
            type    = "slider",
            label   = "SNR Power",
            default = 1.3,
            min     = 0.5,
            max     = 3.0,
            step    = 0.1,
            hint    = "Falloff exponent. Active when Look-Back Smoother is on. Higher=smoother fade at low sigma.",
        },
        {
            key     = "rms_servo",
            type    = "toggle",
            label   = "RMS Servo",
            default = true,
            hint    = "Downward-only energy ceiling. Prevents latent energy accumulation. Off by default — tune min/max for your domain before enabling.",
        },
        {
            key     = "rms_target_min",
            type    = "slider",
            label   = "RMS Target Min",
            default = 1.0,
            min     = 0.1,
            max     = 3.0,
            step    = 0.05,
            hint    = "RMS ceiling at low sigma (late steps). Active when RMS Servo is on. ACE-Step latents ~2.0 RMS. Start at 1.2-1.8.",
        },
        {
            key     = "rms_target_max",
            type    = "slider",
            label   = "RMS Target Max",
            default = 2.2,
            min     = 0.5,
            max     = 4.0,
            step    = 0.05,
            hint    = "RMS ceiling at high sigma (early steps). Active when RMS Servo is on. ACE-Step latents ~2.0 RMS. Start at 2.0-2.5.",
        },
        {
            key     = "rms_servo_gain",
            type    = "slider",
            label   = "Servo Gain",
            default = 0.75,
            min     = 0.1,
            max     = 1.0,
            step    = 0.05,
            hint    = "Servo correction aggressiveness. Active when RMS Servo is on. 0.6=soft, 1.0=hard snap.",
        },
        {
            key     = "seed",
            type    = "slider",
            label   = "Seed",
            default = 42,
            min     = 0,
            max     = 999999,
            step    = 1,
            hint    = "RNG seed for noise generation.",
        },
    },
}

-- ── State (file-level locals, reset when n changes) ──────────────────────────

local _prev_x       = nil   -- for momentum
local _prev_noise   = nil   -- for noise coherence
local _look_back_xp = nil   -- for look-back smoother
local _sigma_max    = nil   -- captured at step 0
local _last_n       = 0

-- Hoisted scratch tables — reused every step to avoid GC pressure
-- Initialized on first step or when n changes
local _noise_buf    = {}    -- reusable noise array
local _x_next_buf   = {}    -- reusable output array
local _x_copy_buf   = {}    -- reusable momentum copy

local EPSILON = 1e-8

-- ── Helpers ──────────────────────────────────────────────────────────────────

local function clamp(v, lo, hi)
    if v < lo then return lo end
    if v > hi then return hi end
    return v
end

-- Seeded LCG RNG — deterministic, no dependency on math.random state
local function make_rng(seed)
    local state = math.floor(seed) % 2147483647
    if state <= 0 then state = state + 2147483646 end
    return function()
        state = (state * 1664525 + 1013904223) % 2147483648
        return state / 2147483648.0
    end
end

-- Box-Muller: two uniform [0,1] → one standard normal sample
local function normal(u1, u2)
    return math.sqrt(-2.0 * math.log(math.max(u1, EPSILON))) * math.cos(2.0 * math.pi * u2)
end

-- Array RMS
local function rms(arr, n)
    local s = 0.0
    for i = 0, n - 1 do s = s + arr[i] * arr[i] end
    return math.sqrt(s / n + EPSILON)
end

-- ── Required step() function ──────────────────────────────────────────────────

function step(xt, vt, t_curr, t_prev, n)
    local step_idx_ = step_index or 0

    -- Reset state on new generation: n change OR step 0 of any new run.
    -- Must check step_idx_==0 because same-length generations won't trigger n change,
    -- causing momentum/look-back to bleed finished audio from the previous run into
    -- the noise of the new one — explosive velocity on step 1.
    if n ~= _last_n or step_idx_ == 0 then
        _prev_x       = nil
        _prev_noise   = nil
        _look_back_xp = nil
        _sigma_max    = nil
        _last_n       = n
        -- Pre-size scratch tables for this n
        for i = 0, n - 1 do
            _noise_buf[i]  = 0.0
            _x_next_buf[i] = 0.0
            _x_copy_buf[i] = 0.0
        end
    end

    -- Read params with safe fallbacks
    local anc_strength   = (params and params.ancestral_strength)   or 1.0
    local noise_coh      = (params and params.noise_coherence)       or 0.0
    local mom_str        = (params and params.momentum_strength)     or 0.1
    local lb_enabled     = (params and params.look_back_enabled)     or false
    local lb_lambda      = (params and params.look_back_lambda)      or 0.55
    local lb_snr_power   = (params and params.look_back_snr_power)   or 1.3
    local rms_servo_on   = (params and params.rms_servo)             or true
    local rms_tgt_min    = (params and params.rms_target_min)        or 1.2
    local rms_tgt_max    = (params and params.rms_target_max)        or 2.2
    local rms_servo_gain = (params and params.rms_servo_gain)        or 0.6
    local seed           = math.floor((params and params.seed)       or 42)

    -- Capture sigma_max on first step for ratio computation
    if _sigma_max == nil then _sigma_max = t_curr end
    local sigma_max   = _sigma_max
    local sigma_ratio = clamp(t_curr / math.max(sigma_max, EPSILON), 0.0, 1.0)

    -- dt = t_prev - t_curr. In flow-matching, t steps DOWN (1→0),
    -- HOT-Step API: t_curr=high sigma, t_prev=lower target. t_curr > t_prev. dt=t_prev-t_curr is NEGATIVE.
    local dt = t_prev - t_curr

    -- Save current xt for momentum (reuse hoisted buffer)
    for i = 0, n - 1 do _x_copy_buf[i] = xt[i] end

    -- ── NOISE GENERATION ──────────────────────────────────────────────────────
    -- Write into hoisted buffer — no table allocation per step
    local rng = make_rng(seed + step_idx_ * 7919)
    for i = 0, n - 1 do
        local u1 = math.max(rng(), EPSILON)
        local u2 = rng()
        _noise_buf[i] = normal(u1, u2)
    end

    -- Noise coherence: blend with carried noise from previous step
    if noise_coh > 0.0 and _prev_noise ~= nil then
        for i = 0, n - 1 do
            _noise_buf[i] = _noise_buf[i] * (1.0 - noise_coh) + _prev_noise[i] * noise_coh
        end
    end
    -- Store for next step — reuse _prev_noise table if same size
    if _prev_noise == nil then _prev_noise = {} end
    for i = 0, n - 1 do _prev_noise[i] = _noise_buf[i] end

    -- ── ANCESTRAL STEP ────────────────────────────────────────────────────────
    -- Variance-preserving SDE noise for flow matching:
    --   noise_scale = sqrt(t_prev^2 - t_curr^2) * anc_strength
    -- t_curr > t_prev, so t_curr^2 - t_prev^2 > 0. Confirmed numerically.
    local noise_scale = math.sqrt(math.max(t_curr * t_curr - t_prev * t_prev, 0.0)) * anc_strength
    if noise_scale > EPSILON then
        for i = 0, n - 1 do
            _x_next_buf[i] = xt[i] + dt * vt[i] + _noise_buf[i] * noise_scale
        end
    else
        for i = 0, n - 1 do _x_next_buf[i] = xt[i] + dt * vt[i] end
    end

    -- ── MOMENTUM ──────────────────────────────────────────────────────────────
    if mom_str > 0.0 and _prev_x ~= nil then
        for i = 0, n - 1 do
            local vel = _x_copy_buf[i] - _prev_x[i]
            _x_next_buf[i] = _x_next_buf[i] + vel * mom_str
        end
    end

    -- ── LOOK-BACK SNR SMOOTHER ────────────────────────────────────────────────
    -- λ(σ) = lb_lambda * (σ/σ_max)^lb_snr_power — heavy early, fades late.
    if lb_enabled then
        local lb_w = lb_lambda * (sigma_ratio ^ lb_snr_power)
        if _look_back_xp == nil then
            _look_back_xp = {}
            for i = 0, n - 1 do
                local u1 = math.max(rng(), EPSILON)
                local u2 = rng()
                _look_back_xp[i] = _x_next_buf[i] + normal(u1, u2) * sigma_max * 0.1
            end
        end
        for i = 0, n - 1 do
            _x_next_buf[i] = _x_next_buf[i] * (1.0 - lb_w) + _look_back_xp[i] * lb_w
        end
        -- Update look-back buffer in-place
        for i = 0, n - 1 do _look_back_xp[i] = _x_next_buf[i] end
    end

    -- ── RMS SERVO (DOWNWARD ONLY) ─────────────────────────────────────────────
    if rms_servo_on then
        local rms_target = rms_tgt_min + (sigma_ratio ^ 0.6) * (rms_tgt_max - rms_tgt_min)
        local cur_rms = rms(_x_next_buf, n)
        if cur_rms > rms_target then
            local servo_rms = cur_rms + rms_servo_gain * (rms_target - cur_rms)
            local scale = servo_rms / cur_rms
            for i = 0, n - 1 do _x_next_buf[i] = _x_next_buf[i] * scale end
        end
    end

    -- ── UPDATE STATE & WRITE OUTPUT ───────────────────────────────────────────
    -- Store momentum reference — reuse table, copy values
    if _prev_x == nil then _prev_x = {} end
    for i = 0, n - 1 do _prev_x[i] = _x_copy_buf[i] end
    for i = 0, n - 1 do xt[i] = _x_next_buf[i] end
end