--[[
storm_sampler_core.lua
STORM -- Stabilized Taylor Oscillation with Runge-Kutta Memory
Adaptive hybrid solver: STORK (stiff) + DPM++3M (stable), per-step dispatch

© 2026 Alexander Allan (MDMAchine) | A&E Concepts
GPL v3 -- Public version. Gradient norm stiffness detection only.

Adapted for HOT-Step full-loop plugin API (owns_loop = true).
All data uses 0-indexed FloatArray or 0-indexed Lua tables.

Version: 3.0.0  (HOT-Step plugin port from v2.1.0)
--]]

solver = {
    name        = "storm",
    display     = "STORM",
    description = "Adaptive STORK/DPM++3M hybrid with stiffness detection",
    accent      = "cyan",
    nfe         = 0,
    order       = 5,
    needs_model = false,
    stateful    = true,
    stochastic  = false,
    owns_loop   = true,
    params      = {
        { key = "stiffness_threshold", type = "slider", label = "Detail Sensitivity",
          default = 0.15, min = 0.05, max = 0.50, step = 0.01,
          hint = "How aggressively complex passages get extra precision. Lower = more careful on transients and busy sections, higher = faster but looser" },
        { key = "look_back_lambda", type = "slider", label = "Coherence Smoothing",
          default = 0.15, min = 0, max = 1, step = 0.01,
          hint = "Blends each step with previous ones for smoother output. 0 = off (raw), higher = more coherent but softer detail" },
        { key = "look_back_snr_power", type = "slider", label = "Early-Step Focus",
          default = 1.5, min = 0.5, max = 3, step = 0.1,
          hint = "Concentrates smoothing on early noisy steps (structure). Higher = smooths structure more, leaves fine detail alone" },
        { key = "rk_order", type = "select", label = "Precision Level",
          default = "auto",
          options = {
            { value = "auto", label = "Auto (Recommended)" },
            { value = "2",    label = "Low (RK2)" },
            { value = "3",    label = "Medium (RK3)" },
            { value = "4",    label = "High (RK4)" },
            { value = "5",    label = "Maximum (RK5)" },
          },
          hint = "Solver accuracy per step. Auto ramps up gradually. Higher = cleaner but more compute per step" },
        { key = "cache_depth", type = "slider", label = "History Memory",
          default = 5, min = 2, max = 10, step = 1,
          hint = "How many previous steps the solver remembers. More = smoother multi-step blending, but diminishing returns past 5" },
        { key = "verbose", type = "toggle", label = "Verbose Logging",
          default = false,
          hint = "Print per-step solver decisions to the console (debug)" },
    },
}

-- ─────────────────────────────────────────────────────────────────────────────
-- HELPERS: FloatArray ↔ Lua table (0-indexed)
-- ─────────────────────────────────────────────────────────────────────────────

local function fa_to_tbl(fa, n)
    local t = {}
    for i = 0, n - 1 do t[i] = fa[i] end
    return t
end

local function tbl_to_fa(t, fa, n)
    for i = 0, n - 1 do fa[i] = t[i] end
end

-- ─────────────────────────────────────────────────────────────────────────────
-- MATH HELPERS (0-indexed Lua tables)
-- ─────────────────────────────────────────────────────────────────────────────

local function vec_norm(v, n)
    local s = 0.0
    for i = 0, n - 1 do s = s + v[i] * v[i] end
    return math.sqrt(s)
end

local function vec_sub_norm(a, b, n)
    local s = 0.0
    for i = 0, n - 1 do local d = a[i] - b[i]; s = s + d * d end
    return math.sqrt(s)
end

local function vec_dot(a, b, n)
    local s = 0.0
    for i = 0, n - 1 do s = s + a[i] * b[i] end
    return s
end

local function vec_clone(v, n)
    local c = {}
    for i = 0, n - 1 do c[i] = v[i] end
    return c
end

local function has_nan_inf_tbl(v, n)
    for i = 0, n - 1 do
        if v[i] ~= v[i] or math.abs(v[i]) == math.huge then return true end
    end
    return false
end

local function clamp(x, lo, hi) return math.max(lo, math.min(hi, x)) end

local function randn_iso(n, scale)
    local out = {}
    for i = 0, n - 1, 2 do
        local u1 = math.max(1e-12, math.random())
        local u2 = math.random()
        local r  = scale * math.sqrt(-2.0 * math.log(u1))
        out[i]   = r * math.cos(2 * math.pi * u2)
        if i + 1 < n then
            out[i + 1] = r * math.sin(2 * math.pi * u2)
        end
    end
    return out
end

-- ─────────────────────────────────────────────────────────────────────────────
-- LOOK-BACK SMOOTHER (arXiv:2602.09449)
-- ─────────────────────────────────────────────────────────────────────────────

local function look_back_smooth(x_curr, x_prev, sigma_curr, sigma_max, lambda_base, snr_power, n)
    if x_prev == nil then return x_curr, 0.0 end
    local ratio = math.min(sigma_curr / math.max(sigma_max, 1e-8), 1.0)
    local lam   = lambda_base * (ratio ^ snr_power)
    local out   = {}
    for i = 0, n - 1 do out[i] = (1.0 - lam) * x_curr[i] + lam * x_prev[i] end
    return out, lam
end

-- ─────────────────────────────────────────────────────────────────────────────
-- STIFFNESS DETECTION
-- ─────────────────────────────────────────────────────────────────────────────

local function compute_stiffness(v_curr, v_cache, step_idx, baseline, threshold, ema_alpha, n_calib, n)
    threshold = threshold or 0.15
    ema_alpha = ema_alpha or 0.3
    n_calib   = n_calib   or 4

    if #v_cache < 1 then return true, baseline, nil end

    local v_prev     = v_cache[#v_cache].v
    local norm_delta = vec_sub_norm(v_curr, v_prev, n)
    local norm_curr  = vec_norm(v_curr, n) + 1e-8
    local raw_ratio  = norm_delta / norm_curr

    local prev_ema = baseline.ema or raw_ratio
    local smoothed = ema_alpha * raw_ratio + (1.0 - ema_alpha) * prev_ema
    baseline.ema   = smoothed

    local dot     = vec_dot(v_curr, v_prev, n)
    local nc      = vec_norm(v_curr, n)
    local np_     = vec_norm(v_prev, n)
    local cos_sim = dot / (nc * np_ + 1e-8)

    if step_idx < n_calib then
        baseline.sum        = (baseline.sum   or 0.0) + smoothed
        baseline.count      = (baseline.count or 0)   + 1
        baseline.last_ratio = smoothed
        return true, baseline, cos_sim
    end

    local bmean    = baseline.sum / math.max(baseline.count, 1)
    local adap_thr = threshold * (bmean / 0.15)
    adap_thr       = clamp(adap_thr, 0.05, 0.50)

    local stiff = smoothed > adap_thr
    baseline.last_ratio     = smoothed
    baseline.last_threshold = adap_thr
    return stiff, baseline, cos_sim
end

-- ─────────────────────────────────────────────────────────────────────────────
-- STORK MULTI-ORDER (AB2-AB5, single NFE, cached derivatives)
-- ─────────────────────────────────────────────────────────────────────────────

local function stork_step(v_cache, x, sigma_curr, sigma_next, v_curr, rk_order, n)
    local dt      = sigma_next - sigma_curr
    local n_cache = #v_cache

    local actual_order
    if rk_order == "auto" then
        actual_order = (n_cache >= 1) and math.min(n_cache + 1, 5) or 1
    else
        actual_order = (n_cache >= 1) and math.min(tonumber(rk_order), n_cache + 1) or 1
    end
    actual_order = math.max(actual_order, 1)

    if n_cache < 1 or actual_order <= 1 then
        local x_next = {}
        for i = 0, n - 1 do x_next[i] = x[i] + dt * v_curr[i] end
        return x_next, 1
    end

    local e0         = v_cache[#v_cache]
    local v_prev_0   = e0.v
    local sigma_prev = e0.sigma

    -- Curvature damping
    local dot     = vec_dot(v_curr, v_prev_0, n)
    local nc      = vec_norm(v_curr, n)
    local np_     = vec_norm(v_prev_0, n)
    local cos_sim = dot / (nc * np_ + 1e-8)
    local damping = clamp(cos_sim, 0.0, 1.0)

    local denom = sigma_curr - sigma_prev
    if math.abs(denom) < 1e-8 then
        local x_next = {}
        for i = 0, n - 1 do x_next[i] = x[i] + dt * v_curr[i] end
        return x_next, 2
    end
    local alpha = (sigma_next - sigma_curr) / denom

    local x_next = {}

    if actual_order == 2 then
        for i = 0, n - 1 do
            local v_extrap = v_curr[i] + (alpha * damping) * (v_curr[i] - v_prev_0[i])
            x_next[i]     = x[i] + dt * (0.5 * v_curr[i] + 0.5 * v_extrap)
        end

    elseif actual_order == 3 and n_cache >= 2 then
        local v1, s1 = v_cache[#v_cache].v,     v_cache[#v_cache].sigma
        local v2, s2 = v_cache[#v_cache - 1].v, v_cache[#v_cache - 1].sigma
        local h  = sigma_curr - s1
        local h1 = s1 - s2
        if math.abs(h) < 1e-8 or math.abs(h1) < 1e-8 then
            for i = 0, n - 1 do
                local ve = v_curr[i] + (alpha * damping) * (v_curr[i] - v1[i])
                x_next[i] = x[i] + dt * (0.5 * v_curr[i] + 0.5 * ve)
            end
            actual_order = 2
        else
            local c0 = 1.0 + (dt / (2.0 * h)) + (dt ^ 2 / (3.0 * h * h1))
            local c1 = -(dt / (2.0 * h)) * (1.0 + dt / h1)
            local c2 = (dt ^ 2) / (3.0 * h * h1)
            for i = 0, n - 1 do
                local v_pred = c0 * v_curr[i] + c1 * v1[i] + c2 * v2[i]
                x_next[i]   = x[i] + dt * (v_curr[i] + damping * (v_pred - v_curr[i]))
            end
        end

    elseif actual_order == 4 and n_cache >= 3 then
        local v1, s1 = v_cache[#v_cache].v,     v_cache[#v_cache].sigma
        local v2, s2 = v_cache[#v_cache - 1].v, v_cache[#v_cache - 1].sigma
        local v3, s3 = v_cache[#v_cache - 2].v, v_cache[#v_cache - 2].sigma
        local h  = sigma_curr - s1
        local h1 = s1 - s2
        local h2 = s2 - s3
        if math.abs(h) < 1e-8 or math.abs(h1) < 1e-8 or math.abs(h2) < 1e-8 then
            local c0 = 1.0 + (dt / (2.0 * h)) + (dt ^ 2 / (3.0 * h * h1))
            local c1 = -(dt / (2.0 * h)) * (1.0 + dt / h1)
            local c2 = (dt ^ 2) / (3.0 * h * h1)
            for i = 0, n - 1 do
                local vp = c0 * v_curr[i] + c1 * v1[i] + c2 * v2[i]
                x_next[i] = x[i] + dt * (v_curr[i] + damping * (vp - v_curr[i]))
            end
            actual_order = 3
        else
            local c0 = 1.0 + dt / (2.0 * h) + dt ^ 2 / (3.0 * h * h1) + dt ^ 3 / (4.0 * h * h1 * h2)
            local c1 = -(dt / (2.0 * h)) * (1.0 + dt / h1 + dt ^ 2 / (2.0 * h1 * h2))
            local c2 = (dt ^ 2 / (3.0 * h * h1)) * (1.0 + dt / (2.0 * h2))
            local c3 = -(dt ^ 3) / (4.0 * h * h1 * h2)
            for i = 0, n - 1 do
                local vp = c0 * v_curr[i] + c1 * v1[i] + c2 * v2[i] + c3 * v3[i]
                x_next[i] = x[i] + dt * (v_curr[i] + damping * (vp - v_curr[i]))
            end
        end

    elseif actual_order >= 5 and n_cache >= 4 then
        local v1, s1 = v_cache[#v_cache].v,     v_cache[#v_cache].sigma
        local v2, s2 = v_cache[#v_cache - 1].v, v_cache[#v_cache - 1].sigma
        local v3, s3 = v_cache[#v_cache - 2].v, v_cache[#v_cache - 2].sigma
        local v4, s4 = v_cache[#v_cache - 3].v, v_cache[#v_cache - 3].sigma
        local h  = sigma_curr - s1
        local h1 = s1 - s2
        local h2 = s2 - s3
        local h3 = s3 - s4
        if math.abs(h) < 1e-8 or math.abs(h1) < 1e-8 or math.abs(h2) < 1e-8 or math.abs(h3) < 1e-8 then
            local c0 = 1.0 + dt / (2.0 * h) + dt ^ 2 / (3.0 * h * h1) + dt ^ 3 / (4.0 * h * h1 * h2)
            local c1 = -(dt / (2.0 * h)) * (1.0 + dt / h1 + dt ^ 2 / (2.0 * h1 * h2))
            local c2 = (dt ^ 2 / (3.0 * h * h1)) * (1.0 + dt / (2.0 * h2))
            local c3 = -(dt ^ 3) / (4.0 * h * h1 * h2)
            for i = 0, n - 1 do
                local vp = c0 * v_curr[i] + c1 * v1[i] + c2 * v2[i] + c3 * v3[i]
                x_next[i] = x[i] + dt * (v_curr[i] + damping * (vp - v_curr[i]))
            end
            actual_order = 4
        else
            local c0 = 1.0 + dt / (2.0 * h) + dt ^ 2 / (3.0 * h * h1) + dt ^ 3 / (4.0 * h * h1 * h2) + dt ^ 4 / (5.0 * h * h1 * h2 * h3)
            local c1 = -(dt / (2.0 * h)) * (1.0 + dt / h1 + dt ^ 2 / (2.0 * h1 * h2) + dt ^ 3 / (3.0 * h1 * h2 * h3))
            local c2 = (dt ^ 2 / (3.0 * h * h1)) * (1.0 + dt / (2.0 * h2) + dt ^ 2 / (3.0 * h2 * h3))
            local c3 = -(dt ^ 3 / (4.0 * h * h1 * h2)) * (1.0 + dt / (2.0 * h3))
            local c4 = dt ^ 4 / (5.0 * h * h1 * h2 * h3)
            for i = 0, n - 1 do
                local vp = c0 * v_curr[i] + c1 * v1[i] + c2 * v2[i] + c3 * v3[i] + c4 * v4[i]
                x_next[i] = x[i] + dt * (v_curr[i] + damping * (vp - v_curr[i]))
            end
            actual_order = 5
        end

    else
        -- Fallback AB2
        for i = 0, n - 1 do
            local ve = v_curr[i] + (alpha * damping) * (v_curr[i] - v_prev_0[i])
            x_next[i] = x[i] + dt * (0.5 * v_curr[i] + 0.5 * ve)
        end
        actual_order = 2
    end

    return x_next, actual_order
end

-- ─────────────────────────────────────────────────────────────────────────────
-- DPM++3M -- smooth schedule path
-- ─────────────────────────────────────────────────────────────────────────────

local function dpmpp3m_step(v_cache, x, sigma_curr, sigma_next, v_curr, n)
    local dt     = sigma_next - sigma_curr
    local x_next = {}

    if #v_cache >= 2 then
        local v1, s1 = v_cache[#v_cache].v,     v_cache[#v_cache].sigma
        local v2, s2 = v_cache[#v_cache - 1].v, v_cache[#v_cache - 1].sigma
        local h  = sigma_curr - s1
        local h1 = s1 - s2
        if math.abs(h) < 1e-8 or math.abs(h1) < 1e-8 then
            for i = 0, n - 1 do x_next[i] = x[i] + dt * v_curr[i] end
        else
            local cc = 1.0 + (dt / (2.0 * h)) + (dt ^ 2 / (3.0 * h * h1))
            local c1 = -(dt / (2.0 * h)) * (1.0 + dt / h1)
            local c2 = (dt ^ 2) / (3.0 * h * h1)
            for i = 0, n - 1 do x_next[i] = x[i] + dt * (cc * v_curr[i] + c1 * v1[i] + c2 * v2[i]) end
        end
    elseif #v_cache >= 1 then
        local v1, s1 = v_cache[#v_cache].v, v_cache[#v_cache].sigma
        local h = sigma_curr - s1
        if math.abs(h) < 1e-8 then
            for i = 0, n - 1 do x_next[i] = x[i] + dt * v_curr[i] end
        else
            for i = 0, n - 1 do x_next[i] = x[i] + dt * (v_curr[i] + (dt / (2.0 * h)) * (v_curr[i] - v1[i])) end
        end
    else
        for i = 0, n - 1 do x_next[i] = x[i] + dt * v_curr[i] end
    end

    return x_next
end

-- ─────────────────────────────────────────────────────────────────────────────
-- SAMPLE — Full-loop entry point
-- ─────────────────────────────────────────────────────────────────────────────

function sample(xt, vt_buf, schedule, n, model_fn)
    -- Read params
    local p = params or {}
    local thr           = p.stiffness_threshold or 0.15
    local lb_lambda     = p.look_back_lambda    or 0.35
    local lb_snr_pow    = p.look_back_snr_power or 1.5
    local rk_order      = p.rk_order            or "auto"
    local depth_max     = p.cache_depth         or 5
    local verbose       = p.verbose             or false

    local hyst      = 0.05
    local ema_a     = 0.3
    local calib_frac = 0.12

    local ns        = #schedule
    -- The engine's schedule table has NO trailing 0 (unlike ComfyUI sigmas):
    -- iterate all ns entries so the last iteration gets sigma_next = 0.0 and
    -- the terminal branch performs the final x0 projection. With ns - 1 the
    -- terminal branch is dead code and the output keeps ~final-sigma noise.
    local n_steps   = ns
    if n_steps < 1 then return end

    local v_cache   = {}
    local baseline  = { sum = 0.0, count = 0 }
    local sigma_max = schedule[1]
    local n_calib   = math.max(2, math.min(5, math.floor(n_steps * calib_frac)))
    local lb_enabled = (lb_lambda > 0)

    if verbose then
        print(string.format("[STORM] Schedule: %d steps | Calib: %d | RK: %s | Cache: %d",
            n_steps, n_calib, tostring(rk_order), depth_max))
    end

    -- Working copy of xt as a Lua table (we write back to FloatArray at each step)
    local x = fa_to_tbl(xt, n)

    -- Seed x_prev for look-back
    local x_prev_lb = nil
    if lb_enabled then
        x_prev_lb = {}
        for i = 0, n - 1 do
            local u1 = math.max(1e-12, math.random())
            local u2 = math.random()
            local r  = (sigma_max * 0.1) * math.sqrt(-2.0 * math.log(u1))
            x_prev_lb[i] = x[i] + r * math.cos(2 * math.pi * u2)
        end
    end

    -- Helper: evaluate model and return velocity as Lua table
    local function eval_model(x_tbl)
        tbl_to_fa(x_tbl, xt, n)
        model_fn(xt, 0)  -- dummy t, we'll set it properly below
        return fa_to_tbl(vt_buf, n)
    end

    -- Proper eval: writes x_tbl to xt, calls model at t_val, returns velocity table
    local function eval_at(x_tbl, t_val)
        tbl_to_fa(x_tbl, xt, n)
        model_fn(xt, t_val)
        return fa_to_tbl(vt_buf, n)
    end

    for i = 1, n_steps do
        local sigma_curr = schedule[i]
        local sigma_next = (i < ns) and schedule[i + 1] or 0.0

        -- Terminal step
        if sigma_next == 0.0 then
            local v_final = eval_at(x, sigma_curr)
            for j = 0, n - 1 do x[j] = x[j] - v_final[j] * sigma_curr end
            if verbose then
                print(string.format("[STORM] Step %02d: FINAL (Euler terminal)", i - 1))
            end
            break
        end

        local x_prev_lb_before = nil
        if lb_enabled then x_prev_lb_before = vec_clone(x, n) end

        -- Evaluate velocity
        local v_curr = eval_at(x, sigma_curr)

        -- Stiffness detection
        local stiff, cos_sim_out
        if #v_cache >= 1 then
            stiff, baseline, cos_sim_out = compute_stiffness(
                v_curr, v_cache, i - 1, baseline, thr, ema_a, n_calib, n)
        else
            stiff, cos_sim_out = true, nil
        end

        -- Hysteresis
        local prev_mode = baseline.prev_mode or "STORK"
        if prev_mode == "DPM++" and not stiff then
            if (baseline.last_ratio or 0) > (baseline.last_threshold or thr) + hyst then
                stiff = true
            end
        end

        -- Dispatch
        local x_next, actual_order, mode
        if stiff then
            x_next, actual_order = stork_step(v_cache, x, sigma_curr, sigma_next, v_curr, rk_order, n)
            mode = "STORK"
        else
            x_next = dpmpp3m_step(v_cache, x, sigma_curr, sigma_next, v_curr, n)
            mode         = "DPM++"
            actual_order = 3
        end

        -- Verbose
        if verbose then
            local lr  = baseline.last_ratio     or 0.0
            local lt  = baseline.last_threshold or thr
            local cs  = cos_sim_out and string.format("%.4f", cos_sim_out) or "N/A"
            local tag = (stiff and prev_mode == "DPM++") and " -> CURVATURE SPIKE" or ""
            print(string.format("[STORM] Step %02d: %-5s RK%d | Ratio: %.3f | Thr: %.3f | cos: %s%s",
                i - 1, mode, actual_order, lr, lt, cs, tag))
        end

        -- NaN guard
        if has_nan_inf_tbl(x_next, n) then
            print(string.format("[STORM] NaN/Inf at step %d. Flushing cache.", i - 1))
            local dt = sigma_next - sigma_curr
            x_next = {}
            for j = 0, n - 1 do x_next[j] = x[j] + dt * v_curr[j] end
            v_cache = {}
            baseline.prev_mode = "STORK"
            actual_order = 1
        end

        -- Update cache
        table.insert(v_cache, { v = v_curr, sigma = sigma_curr })
        while #v_cache > depth_max do table.remove(v_cache, 1) end
        baseline.prev_mode = mode

        x = x_next

        -- Look-Back smoothing
        if lb_enabled and x_prev_lb ~= nil then
            local lam
            x, lam = look_back_smooth(x, x_prev_lb, sigma_curr, sigma_max, lb_lambda, lb_snr_pow, n)
            if verbose then
                print(string.format("[STORM] LookBack λ=%.4f @ σ=%.3f", lam, sigma_curr))
            end
        end
        x_prev_lb = x_prev_lb_before

        -- Write x back to xt FloatArray for on_step hooks (DCW, repaint)
        tbl_to_fa(x, xt, n)
        -- Write velocity to vt_buf for DCW
        tbl_to_fa(v_curr, vt_buf, n)

        -- Report step (engine hooks: DCW, repaint, progress)
        if on_step(i - 1, sigma_curr, sigma_next) then return end

        -- Re-read xt in case hooks modified it (DCW, repaint)
        x = fa_to_tbl(xt, n)
    end

    -- Write final x0 to xt
    tbl_to_fa(x, xt, n)
end