-- ============================================================================
-- 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 Causal Scheduler v2.1 — LINA Time Warping + Multi-Mode Base Curves
-- MDMAchine | A&E Concepts © 2026
--
-- Port of md_causal_scheduler_core.py to HOT-Step-CPP Lua.
--
-- 14 BASE SCHEDULE MODES:
--   karras     — power-law rho spacing (rho=7 default, Karras et al.)
--   simple     — smoothstep (cubic hermite: t*t*(3-2t))
--   linear     — uniform spacing
--   exponential — exp decay: sigma_max * (sigma_min/sigma_max)^t
--   polynomial  — power curve: linspace of sigma^(1/power)
--   beta        — beta distribution curve (alpha, beta params)
--   ays         — AYS adaptive schedule (sigmoid + concentration blend)
--   bong        — tangent-based 2-phase schedule (pivot point)
--   linear_quadratic — linear phase then quadratic phase
--   ddim_uniform — DDIM-style uniform timestep mapping
--   sgm_uniform  — SGM uniform (linear 999→0 mapped to sigma range)
--   blended     — karras + linear blend by blend_factor
--   variance_preserving — log-space interpolation
--   kl_optimal  — arctan-based KL-optimal spacing
--
-- LINA WARP:
--   Post-processes any base schedule by warping the time axis:
--   t_warped = t^shift (power warp on the CDF index)
--   shift < 1: front-loads steps (more at high sigma)
--   shift > 1: back-loads steps (more at low sigma)
--   shift = 1: no warp (identity)
-- ============================================================================

scheduler = {
    name        = "md_causal",
    display     = "MD Causal (LINA + 14 Modes)",
    description = "14 base schedule modes with LINA time-axis warp. Karras, smoothstep, beta, AYS, bong, DDIM, SGM, blended, variance-preserving, KL-optimal and more. Port of md_causal_scheduler_core v2.1.",
    params      = {
        {
            key     = "mode",
            type    = "select",
            label   = "Schedule Mode",
            default = "polynomial",
            options = {
                { value = "karras",              label = "Karras (rho)"          },
                { value = "simple",              label = "Simple (Smoothstep)"   },
                { value = "linear",              label = "Linear"                },
                { value = "exponential",         label = "Exponential"           },
                { value = "polynomial",          label = "Polynomial"            },
                { value = "beta",                label = "Beta"                  },
                { value = "ays",                 label = "AYS"                   },
                { value = "bong",                label = "Bong (Tangent)"        },
                { value = "linear_quadratic",    label = "Linear-Quadratic"      },
                { value = "ddim_uniform",        label = "DDIM Uniform"          },
                { value = "sgm_uniform",         label = "SGM Uniform"           },
                { value = "blended",             label = "Blended (Karras+Lin)"  },
                { value = "variance_preserving", label = "Variance Preserving"   },
                { value = "kl_optimal",          label = "KL Optimal"            },
            },
            hint = "Base schedule curve before LINA warp is applied.",
        },
        {
            key     = "lina_shift",
            type    = "slider",
            label   = "LINA Shift",
            default = 1.2,
            min     = 0.1,
            max     = 3.0,
            step    = 0.05,
            hint    = "Time-axis warp. 1.0=none. <1=front-load (more high-sigma steps). >1=back-load (more low-sigma steps).",
        },
        {
            key     = "rho",
            type    = "slider",
            label   = "Rho (Karras)",
            default = 7.0,
            min     = 1.0,
            max     = 15.0,
            step    = 0.5,
            hint    = "Karras rho parameter. 7=default. Higher=more steps at low sigma.",
            visible_when = { key = "mode", equals = "karras" },
        },
        {
            key     = "power",
            type    = "slider",
            label   = "Power (Polynomial)",
            default = 2.0,
            min     = 0.5,
            max     = 5.0,
            step    = 0.1,
            hint    = "Polynomial exponent. 2=quadratic, 1=linear.",
            visible_when = { key = "mode", equals = "polynomial" },
        },
        {
            key     = "beta_alpha",
            type    = "slider",
            label   = "Beta Alpha",
            default = 0.6,
            min     = 0.1,
            max     = 3.0,
            step    = 0.1,
            hint    = "Beta distribution alpha parameter.",
            visible_when = { key = "mode", equals = "beta" },
        },
        {
            key     = "beta_beta",
            type    = "slider",
            label   = "Beta Beta",
            default = 0.6,
            min     = 0.1,
            max     = 3.0,
            step    = 0.1,
            hint    = "Beta distribution beta parameter.",
            visible_when = { key = "mode", equals = "beta" },
        },
        {
            key     = "blend_factor",
            type    = "slider",
            label   = "Blend Factor",
            default = 0.5,
            min     = 0.0,
            max     = 1.0,
            step    = 0.05,
            hint    = "Blend between Karras (0) and Linear (1).",
            visible_when = { key = "mode", equals = "blended" },
        },
        {
            key     = "bong_pivot",
            type    = "slider",
            label   = "Bong Pivot",
            default = 0.5,
            min     = 0.1,
            max     = 0.9,
            step    = 0.05,
            hint    = "Bong: fraction of steps in compression phase.",
            visible_when = { key = "mode", equals = "bong" },
        },
        {
            key     = "bong_slope_comp",
            type    = "slider",
            label   = "Bong Slope Comp",
            default = 1.2,
            min     = 0.1,
            max     = 3.0,
            step    = 0.1,
            hint    = "Bong: tangent slope in compression phase.",
            visible_when = { key = "mode", equals = "bong" },
        },
        {
            key     = "bong_slope_detail",
            type    = "slider",
            label   = "Bong Slope Detail",
            default = 0.8,
            min     = 0.1,
            max     = 3.0,
            step    = 0.1,
            hint    = "Bong: tangent slope in detail phase.",
            visible_when = { key = "mode", equals = "bong" },
        },
    },
}

local EPSILON = 1e-6
local MONOTONIC_DECAY = 0.99

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

-- ── Base schedule generators ─────────────────────────────────────────────────

local function karras(n, s_min, s_max, rho)
    -- sigma[i] = (s_max^(1/rho) + i/(n-1) * (s_min^(1/rho) - s_max^(1/rho)))^rho
    local inv_rho = 1.0 / rho
    local max_inv = s_max ^ inv_rho
    local min_inv = s_min ^ inv_rho
    local s = {}
    for i = 0, n do
        local t = i / n
        s[i] = (max_inv + t * (min_inv - max_inv)) ^ rho
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function simple(n, s_min, s_max)
    local s = {}
    for i = 0, n do
        local t = i / n
        local smooth = t * t * (3.0 - 2.0 * t)
        s[i] = s_max - (s_max - s_min) * smooth
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function linear(n, s_min, s_max)
    local s = {}
    for i = 0, n do
        s[i] = s_max - (s_max - s_min) * (i / n)
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function exponential(n, s_min, s_max)
    local s = {}
    local safe_max = math.max(s_max, 1e-9)
    for i = 0, n do
        local t = i / n
        s[i] = safe_max * (s_min / safe_max) ^ t
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function polynomial(n, s_min, s_max, power)
    local s = {}
    local inv_p = 1.0 / math.max(power, 0.1)
    local lo = s_min ^ inv_p
    local hi = s_max ^ inv_p
    for i = 0, n do
        local t = i / n
        s[i] = (hi + t * (lo - hi)) ^ power
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function beta_sched(n, s_min, s_max, alpha, beta_)
    local s = {}
    for i = 0, n do
        local t = i / n
        local alpha_ = math.max(alpha, 0.1)
        local beta__ = math.max(beta_, 0.1)
        local beta_curve = clamp(1.0 - (1.0 - t ^ alpha_) ^ beta__, 0.0, 1.0)
        s[i] = s_max * (1.0 - beta_curve) + s_min * beta_curve
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function ays_sched(n, s_min, s_max)
    local s = {}
    for i = 0, n do
        local t = i / n
        -- sigmoid centered at 0.5, steepness 10
        local sig = 1.0 / (1.0 + math.exp(-10.0 * (t - 0.5)))
        -- AYS blend: sigmoid 0.7 + concentration (exp decay) 0.3
        local conc = math.exp(-2.0 * t)
        local ays = sig * 0.7 + conc * 0.3
        -- normalize and invert: high sigma at start
        s[i] = s_min + (s_max - s_min) * (1.0 - ays)
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function bong_sched(n, s_min, s_max, pivot, slope_comp, slope_det)
    local comp_steps = math.max(1, math.floor(n * pivot))
    local det_steps  = math.max(1, n - comp_steps)

    local sigmas = {}
    local pi_half = math.pi / 2.0 - 0.1

    -- Compression phase
    for i = 0, comp_steps - 1 do
        local t = i / math.max(comp_steps - 1, 1)
        local angle = t * pi_half * slope_comp
        local warped = math.tan(angle) / math.tan(pi_half * slope_comp)
        sigmas[i] = s_max * (1.0 - warped * pivot)
    end

    -- Detail phase
    for i = 0, det_steps - 1 do
        local t = i / math.max(det_steps - 1, 1)
        local angle = t * pi_half * slope_det
        local warped = math.tan(angle) / math.tan(pi_half * slope_det)
        local start = s_max * (1.0 - pivot)
        sigmas[comp_steps + i] = start * (1.0 - warped) + s_min * warped
    end

    sigmas[n] = s_min

    -- Enforce monotonic
    for i = 0, n - 1 do
        if sigmas[i] ~= nil and sigmas[i + 1] ~= nil then
            if sigmas[i] <= sigmas[i + 1] then
                sigmas[i + 1] = math.max(sigmas[i] * MONOTONIC_DECAY, sigmas[i] - EPSILON)
            end
        end
    end

    sigmas[0] = s_max; sigmas[n] = s_min
    return sigmas
end

local function ddim_uniform(n, s_min, s_max)
    local max_ts = 1000
    local s = {}
    for i = 0, n do
        local ts = max_ts - i * (max_ts / n)
        s[i] = s_min + (s_max - s_min) * ((ts / max_ts) ^ 0.5)
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function sgm_uniform(n, s_min, s_max)
    local s = {}
    for i = 0, n do
        local t = i / n
        s[i] = s_min + (s_max - s_min) * (1.0 - t)
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function blended(n, s_min, s_max, rho, blend)
    local k = karras(n, s_min, s_max, rho)
    local l = linear(n, s_min, s_max)
    local s = {}
    for i = 0, n do
        s[i] = (1.0 - blend) * k[i] + blend * l[i]
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function variance_preserving(n, s_min, s_max)
    local s = {}
    local log_min = math.log(math.max(s_min, 1e-9))
    local log_max = math.log(math.max(s_max, 1e-9))
    for i = 0, n do
        local t = i / n
        s[i] = math.exp((1.0 - t) * log_max + t * log_min)
    end
    s[0] = s_max; s[n] = s_min
    return s
end

local function kl_optimal(n, s_min, s_max)
    local s = {}
    local atan_min = math.atan(s_min)
    local atan_max = math.atan(s_max)
    for i = 0, n do
        local t = i / n
        s[i] = math.tan((1.0 - t) * atan_max + t * atan_min)
    end
    s[0] = s_max; s[n] = s_min
    return s
end

-- ── LINA warp ─────────────────────────────────────────────────────────────────

local function apply_lina_warp(sigmas, n, shift)
    if shift == 1.0 then return sigmas end
    local warped = {}
    for i = 0, n do
        local t = i / n
        -- Warp: t_warped = t^shift → index into sigma array
        local t_w = t ^ shift
        local raw_idx = t_w * n
        local idx_lo = clamp(math.floor(raw_idx), 0, n)
        local idx_hi = clamp(idx_lo + 1, 0, n)
        local frac   = raw_idx - idx_lo
        local s_lo   = sigmas[idx_lo] or sigmas[n]
        local s_hi   = sigmas[idx_hi] or sigmas[n]
        warped[i] = s_lo * (1.0 - frac) + s_hi * frac
    end
    warped[0] = sigmas[0]
    warped[n] = sigmas[n]
    return warped
end

-- ── Required schedule() function ─────────────────────────────────────────────

function schedule(output, num_steps, shift)
    local mode        = (params and params.mode)         or "karras"
    local lina_shift  = (params and params.lina_shift)   or 1.0
    local rho         = (params and params.rho)          or 7.0
    local power       = (params and params.power)        or 2.0
    local ba          = (params and params.beta_alpha)   or 0.6
    local bb          = (params and params.beta_beta)    or 0.6
    local blend       = (params and params.blend_factor) or 0.5
    local b_pivot     = (params and params.bong_pivot)   or 0.5
    local b_comp      = (params and params.bong_slope_comp)   or 1.2
    local b_det       = (params and params.bong_slope_detail) or 0.8

    local s_max = 1.0
    local s_min = 0.0

    local sigmas
    if     mode == "karras"              then sigmas = karras(num_steps, s_min, s_max, rho)
    elseif mode == "simple"              then sigmas = simple(num_steps, s_min, s_max)
    elseif mode == "linear"              then sigmas = linear(num_steps, s_min, s_max)
    elseif mode == "exponential"         then sigmas = exponential(num_steps, s_min, s_max)
    elseif mode == "polynomial"          then sigmas = polynomial(num_steps, s_min, s_max, power)
    elseif mode == "beta"                then sigmas = beta_sched(num_steps, s_min, s_max, ba, bb)
    elseif mode == "ays"                 then sigmas = ays_sched(num_steps, s_min, s_max)
    elseif mode == "bong"                then sigmas = bong_sched(num_steps, s_min, s_max, b_pivot, b_comp, b_det)
    elseif mode == "ddim_uniform"        then sigmas = ddim_uniform(num_steps, s_min, s_max)
    elseif mode == "sgm_uniform"         then sigmas = sgm_uniform(num_steps, s_min, s_max)
    elseif mode == "blended"             then sigmas = blended(num_steps, s_min, s_max, rho, blend)
    elseif mode == "variance_preserving" then sigmas = variance_preserving(num_steps, s_min, s_max)
    elseif mode == "kl_optimal"          then sigmas = kl_optimal(num_steps, s_min, s_max)
    else                                      sigmas = karras(num_steps, s_min, s_max, rho)
    end

    -- LINA warp
    if lina_shift ~= 1.0 then
        sigmas = apply_lina_warp(sigmas, num_steps, lina_shift)
    end

    -- Native shift warp
    if shift ~= 1.0 then
        for i = 0, num_steps do
            local t = sigmas[i]
            sigmas[i] = shift * t / (1.0 + (shift - 1.0) * t)
        end
    end

    for i = 0, num_steps - 1 do
        output[i] = sigmas[i] or 1.0 - i / num_steps
    end
end