helena 0.1.0

use serde::{Deserialize, Serialize};

use crate::{Error, Result};

use super::NoiseLevel;

/// A per-noise-level schedule for the classifier-free guidance weight.
///
/// Sampling mixes the conditional and unconditional velocity as
/// `v = v_uncond + w·(v_cond − v_uncond)`, where the run's `guidance_scale` sets
/// `w`. Holding `w` constant over the whole chain over-guides a scarce-data prior:
/// the evidence is that guidance is harmful at the highest noise, unnecessary at
/// the lowest, and useful only in the mid band, and that constant strong guidance
/// collapses seed diversity (the scarce-data objective brief, lever H4;
/// Kynkäänniemi et al. 2024, Wang et al. 2024).
///
/// A schedule maps a [`NoiseLevel`] and the run's `w` to the guidance applied at
/// that step. It only ever attenuates `w` toward the neutral pass-through `1.0`
/// (pure conditional), never amplifies past `w`, so the result always lies on the
/// segment between `1.0` and `w`. The sampler therefore needs its unconditional
/// pass on exactly the condition it did before scheduling (`w != 1.0`), and
/// [`Constant`](Self::Constant) reproduces the pre-scheduling velocity bit for bit.
///
/// Total and deterministic: a pure function of the level and `w`, no Burn, no RNG.
#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "GuidanceScheduleRepr", into = "GuidanceScheduleRepr")]
#[non_exhaustive]
pub enum GuidanceSchedule {
    /// Apply the run's guidance at every level: the pre-scheduling behavior.
    #[default]
    Constant,
    /// Full guidance inside an interior signal-variance band, neutral pass-through
    /// outside. Limited-interval CFG (Kynkäänniemi et al. 2024).
    Interval(GuidanceInterval),
    /// Ramp guidance linearly in signal variance from neutral at the highest noise
    /// (`gamma = 0`) to the run's guidance at the lowest noise (`gamma = 1`).
    /// Monotone-increasing CFG (Wang et al. 2024).
    Monotone,
}

impl GuidanceSchedule {
    /// The guidance applied at `level` given the run's base `guidance`.
    ///
    /// [`Constant`](Self::Constant) returns `guidance` unchanged (bit for bit);
    /// [`Interval`](Self::Interval) returns `guidance` inside the band and the
    /// neutral `1.0` outside; [`Monotone`](Self::Monotone) returns the convex blend
    /// `1 + (guidance − 1)·gamma`. Every arm yields a point on the segment between
    /// `1.0` and `guidance`, so scheduling never amplifies past the run's scale.
    pub fn guidance_at(&self, level: NoiseLevel, guidance: f32) -> f32 {
        match self {
            // Return the base scale unchanged so an unscheduled run stays exactly
            // the pre-scheduling sampler, not a round-tripped approximation of it.
            Self::Constant => guidance,
            Self::Interval(band) => {
                if band.contains(level.signal_variance()) {
                    guidance
                } else {
                    1.0
                }
            }
            Self::Monotone => 1.0 + (guidance - 1.0) * level.signal_variance(),
        }
    }
}

/// An interior signal-variance band over which limited-interval CFG keeps full
/// guidance; outside it, guidance passes through to neutral.
///
/// Bounds are signal variances (`gamma`), so both lie in `[0, 1]` and `lo < hi`.
/// The strict order makes a degenerate (empty or point) band unrepresentable, so
/// an [`Interval`](GuidanceSchedule::Interval) schedule always covers a real band.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct GuidanceInterval {
    lo: f32,
    hi: f32,
}

impl GuidanceInterval {
    /// Build a band, requiring `lo` and `hi` finite signal variances in `[0, 1]`
    /// with `lo < hi`.
    pub fn new(lo: f32, hi: f32) -> Result<Self> {
        for (name, value) in [("lo", lo), ("hi", hi)] {
            if !value.is_finite() || !(0.0..=1.0).contains(&value) {
                return Err(Error::validation(format!(
                    "guidance interval {name} must be a finite signal variance in [0, 1], got {value}"
                )));
            }
        }
        if lo >= hi {
            return Err(Error::validation(format!(
                "guidance interval requires lo < hi, got [{lo}, {hi}]"
            )));
        }
        Ok(Self { lo, hi })
    }

    /// The lower signal-variance bound.
    pub fn lo(&self) -> f32 {
        self.lo
    }

    /// The upper signal-variance bound.
    pub fn hi(&self) -> f32 {
        self.hi
    }

    /// Whether `gamma` falls inside the closed band.
    fn contains(&self, gamma: f32) -> bool {
        (self.lo..=self.hi).contains(&gamma)
    }
}

/// Deserialization mirror routed through [`GuidanceInterval::new`] and the
/// [`GuidanceSchedule`] variants.
///
/// Adjacent tagging keeps a flat, self-describing config encoding
/// (`{"type": "interval", "band": {"lo": .., "hi": ..}}`). `deny_unknown_fields`
/// here rejects a stray key beside `type`/`band` (e.g. on a unit variant), and the
/// inner [`GuidanceIntervalRepr`] does the same for the band, so the whole object
/// fails closed at every depth rather than silently dropping a typo.
#[derive(Clone, Serialize, Deserialize)]
#[serde(
    tag = "type",
    content = "band",
    rename_all = "snake_case",
    deny_unknown_fields
)]
enum GuidanceScheduleRepr {
    Constant,
    Interval(GuidanceIntervalRepr),
    Monotone,
}

#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct GuidanceIntervalRepr {
    lo: f32,
    hi: f32,
}

impl TryFrom<GuidanceScheduleRepr> for GuidanceSchedule {
    type Error = Error;

    fn try_from(repr: GuidanceScheduleRepr) -> Result<Self> {
        Ok(match repr {
            GuidanceScheduleRepr::Constant => Self::Constant,
            GuidanceScheduleRepr::Interval(band) => {
                Self::Interval(GuidanceInterval::new(band.lo, band.hi)?)
            }
            GuidanceScheduleRepr::Monotone => Self::Monotone,
        })
    }
}

impl From<GuidanceSchedule> for GuidanceScheduleRepr {
    fn from(schedule: GuidanceSchedule) -> Self {
        match schedule {
            GuidanceSchedule::Constant => Self::Constant,
            GuidanceSchedule::Interval(band) => Self::Interval(GuidanceIntervalRepr {
                lo: band.lo(),
                hi: band.hi(),
            }),
            GuidanceSchedule::Monotone => Self::Monotone,
        }
    }
}