finquant 0.0.59

//! Time-dependent FX-SABR calibration — van der Stoep, Grzelak,
//! Oosterlee (2015) §5.1, Algorithm 1.
//!
//! # Four-stage flow
//!
//! Inputs are a set of `N` pillars `{(Tᵢ, Fᵢ, strikes, market_vols)}`.
//! Output is a time-dependent SABR schedule `(α(t), ρ(t), ν(t))` plus a
//! fixed `β` that matches every pillar's smile.
//!
//! 1. **Market-effective extraction** — at each expiry `Tᵢ`
//!    independently, fit a constant-parameter Hagan SABR to the smile
//!    via the existing [`calibrate_sabr`][csabr]. Produces per-pillar
//!    `(α̃ᵢ, ρ̃ᵢ, ν̃ᵢ)`.
//! 2. **γ(t) sequential** — grow the `ν`-schedule one segment at a time
//!    so [`effective_vol_vol`] at each knot equals the pillar's `ν̃ᵢ`.
//!    Uses the paper's approximation `ω(t) ≈ ω̃ᵢ_mar` for the γ̃
//!    mapping at this stage (paper eq. 42).
//! 3. **ω(t) sequential** — grow the `α`-schedule so
//!    [`effective_term_structure`] matches `α̃ᵢ`.
//! 4. **ρ(t) sequential** — grow the `ρ`-schedule so
//!    [`effective_correlation`] matches `ρ̃ᵢ`. Because the mapping is
//!    linear in `ρ`, each sub-step is a closed-form algebraic solve
//!    (no root find needed).
//!
//! `β` is held fixed throughout (FX convention: `0.5`). The local-vol
//! compensator from paper Phase 5 is out of scope here — see
//! `FX_SABR_PLAN.md`.
//!
//! [csabr]: crate::models::forex::sabr_calibrator::calibrate

use crate::error::Result;
use crate::math::optimize::NelderMeadOptions;
use crate::models::common::calibration::{Calibration, CalibrationReport};
use crate::models::forex::market_data::MarketSmileStrip;
use crate::models::forex::sabr::SabrParams;
use crate::models::forex::sabr_calibrator::{
    CalibrationTarget, calibrate as calibrate_sabr, targets_from_grid,
};
use crate::models::forex::sabr_effective::{
    PiecewiseConstant, effective_correlation, effective_term_structure, effective_vol_vol,
};
use crate::models::forex::sabr_time_dependent::TimeDependentSabrParams;

/// Market data for one expiry.
#[derive(Clone, Debug)]
pub struct PillarTarget {
    pub expiry: f64,
    pub forward: f64,
    pub strikes: Vec<f64>,
    pub market_vols: Vec<f64>,
}

/// Calibration diagnostics per pillar.
#[derive(Clone, Debug)]
pub struct PillarDiagnostics {
    pub expiry: f64,
    pub market_effective: SabrParams,
    /// RMSE of the stage-1 constant-SABR fit (decimals of vol).
    pub stage1_rmse: f64,
    /// RMSE between model-effective and market-effective after all
    /// stages (decimals of vol-equivalent). Captures residual
    /// calibration error from the approximate stage-2 mapping.
    pub post_fit_rmse: f64,
}

#[derive(Clone, Debug)]
pub struct TimeDependentCalibrationResult {
    pub params: TimeDependentSabrParams,
    pub per_pillar: Vec<PillarDiagnostics>,
}

/// Drive the full 4-stage calibration.
///
/// * `pillars` — strictly-increasing `expiry` ordering required.
/// * `beta` — fixed CEV exponent, usually `0.5` for FX.
/// * `forward_0` — `F(0)` for the time-dependent simulator; in
///   practice callers pass the 1 st-pillar forward (the schedule
///   absorbs any rate-scaling differences between pillars via `ω`).
/// * `stage1_options` — Nelder-Mead tuning for the per-pillar fits
///   (identical meaning to [`calibrate_sabr`][csabr]'s `options`).
///
/// [csabr]: crate::models::forex::sabr_calibrator::calibrate
pub fn calibrate_time_dependent(
    pillars: &[PillarTarget],
    beta: f64,
    forward_0: f64,
    stage1_options: NelderMeadOptions,
) -> TimeDependentCalibrationResult {
    assert!(!pillars.is_empty(), "need at least one pillar");
    for w in pillars.windows(2) {
        assert!(
            w[1].expiry > w[0].expiry,
            "pillars must be sorted by expiry"
        );
    }

    // --- Stage 1: per-pillar constant-SABR fits ---------------------
    let market: Vec<SabrParams> = pillars
        .iter()
        .map(|pi| {
            let seed = SabrParams::new(pi.market_vols[pi.market_vols.len() / 2], beta, -0.20, 0.30);
            let tgts: Vec<CalibrationTarget> = targets_from_grid(&pi.strikes, &pi.market_vols);
            let res = calibrate_sabr(seed, pi.forward, &tgts, pi.expiry, stage1_options);
            res.params
        })
        .collect();
    let stage1_rmse: Vec<f64> = pillars
        .iter()
        .zip(&market)
        .map(|(pi, mp)| {
            use crate::models::forex::sabr::hagan_implied_vol;
            let mut ss = 0.0;
            for (k, v) in pi.strikes.iter().zip(&pi.market_vols) {
                let m = hagan_implied_vol(mp, pi.forward, *k, pi.expiry);
                ss += (m - v).powi(2);
            }
            (ss / pi.strikes.len() as f64).sqrt()
        })
        .collect();

    // Knot grid: 0 plus each pillar expiry.
    let knots: Vec<f64> = std::iter::once(0.0)
        .chain(pillars.iter().map(|p| p.expiry))
        .collect();

    // --- Stage 2: grow ν(t) (= "γ" in paper) ------------------------
    //
    // Approximation ω(t) ≈ ω̃ᵢ_mar (paper eq. 42). We iteratively
    // build a ν schedule by bracketing each new segment's value such
    // that `effective_vol_vol` at the running expiry hits ν̃ᵢ.
    let mut nu_values: Vec<f64> = Vec::with_capacity(pillars.len());
    for i in 0..pillars.len() {
        let target = market[i].nu;
        // Candidate "trial" ν on segment i; ω schedule frozen at the
        // market-effective value at this pillar.
        let omega_flat = market[i].alpha;
        let omega_sched = PiecewiseConstant::new(knots[..=i + 1].to_vec(), vec![omega_flat; i + 1]);

        // Build partial ν with the first i segments already fixed;
        // bracket the last segment.
        let partial_nu = |x: f64| -> PiecewiseConstant {
            let mut vs = nu_values.clone();
            vs.push(x);
            PiecewiseConstant::new(knots[..=i + 1].to_vec(), vs)
        };
        let f = |x: f64| -> f64 {
            let sched = partial_nu(x);
            effective_vol_vol(&sched, &omega_sched, pillars[i].expiry) - target
        };

        // Bracket. Try widening if needed. `f` is monotone-increasing
        // in the last segment's ν (more vol-of-vol late in the
        // schedule ⇒ larger effective γ̃).
        let lo = 1.0e-6_f64;
        let mut hi = (target * 2.0).max(0.1);
        let mut f_hi = f(hi);
        let mut bumps = 0;
        while f_hi < 0.0 && bumps < 20 {
            hi *= 2.0;
            f_hi = f(hi);
            bumps += 1;
        }
        assert!(
            f_hi >= 0.0,
            "stage-2 ν bracket for pillar {}: could not reach target ν̃ = {}",
            i,
            target
        );
        let nu_i = bisect(f, lo, hi, 1.0e-10, 80);
        nu_values.push(nu_i);
    }
    let nu_sched = PiecewiseConstant::new(knots.clone(), nu_values.clone());

    // --- Stage 3: grow ω(t) (= "α" in paper) ------------------------
    //
    // Given ν(t), invert the effective-term-structure mapping to
    // recover ω segment-by-segment. For each pillar i, find αᵢ such
    // that `effective_term_structure(ν, α, Tᵢ)` equals α̃ᵢ_mar.
    let mut alpha_values: Vec<f64> = Vec::with_capacity(pillars.len());
    for i in 0..pillars.len() {
        let target = market[i].alpha;
        let partial_alpha = |x: f64| -> PiecewiseConstant {
            let mut vs = alpha_values.clone();
            vs.push(x);
            PiecewiseConstant::new(knots[..=i + 1].to_vec(), vs)
        };
        let sub_nu = PiecewiseConstant::new(knots[..=i + 1].to_vec(), nu_values[..=i].to_vec());
        let f = |x: f64| -> f64 {
            let sched = partial_alpha(x);
            effective_term_structure(&sub_nu, &sched, pillars[i].expiry) - target
        };
        // ω̃² is monotone in ω_i² (linear in ω_i² through I(Tᵢ)), so
        // monotone in ω_i.
        let lo = 1.0e-6_f64;
        let mut hi = (target * 4.0).max(0.5);
        let mut f_hi = f(hi);
        let mut bumps = 0;
        while f_hi < 0.0 && bumps < 20 {
            hi *= 2.0;
            f_hi = f(hi);
            bumps += 1;
        }
        assert!(
            f_hi >= 0.0,
            "stage-3 α bracket for pillar {}: could not reach target α̃ = {}",
            i,
            target
        );
        let alpha_i = bisect(f, lo, hi, 1.0e-10, 80);
        alpha_values.push(alpha_i);
    }
    let alpha_sched = PiecewiseConstant::new(knots.clone(), alpha_values.clone());

    // --- Stage 4: grow ρ(t) -----------------------------------------
    //
    // Lemma 4.6: `ρ̃ = (ω̃ / (γ̃ · Tᵢ)) · ∫(ρ·ν/ω) dt`. The integral
    // is linear in each `ρₖ` so we invert algebraically.
    //
    // Let `S_k = ∫_{knots[k]}^{knots[k+1]} ν(t)/ω(t) dt`. Then
    //   ρ̃ᵢ · (γ̃ · Tᵢ / ω̃) = Σ_{k<i} ρₖ · S_k + ρᵢ · S_i
    // and ρᵢ = (ρ̃ᵢ · (γ̃·T/ω̃) − Σ_{k<i} ρₖ·S_k) / S_i.
    let mut rho_values: Vec<f64> = Vec::with_capacity(pillars.len());
    let mut tail_sum = 0.0_f64;
    for i in 0..pillars.len() {
        let t = pillars[i].expiry;
        let sub_nu = PiecewiseConstant::new(knots[..=i + 1].to_vec(), nu_values[..=i].to_vec());
        let sub_alpha =
            PiecewiseConstant::new(knots[..=i + 1].to_vec(), alpha_values[..=i].to_vec());
        let gamma_tilde = effective_vol_vol(&sub_nu, &sub_alpha, t);
        let omega_tilde = effective_term_structure(&sub_nu, &sub_alpha, t);
        let target = market[i].rho;
        let s_i = (knots[i + 1] - knots[i]) * nu_values[i] / alpha_values[i];
        assert!(
            s_i.abs() > 1.0e-15,
            "stage-4 ρ denominator zero at pillar {}",
            i
        );
        let rhs = target * gamma_tilde * t / omega_tilde;
        let rho_i = (rhs - tail_sum) / s_i;
        tail_sum += rho_i * s_i;

        // Clamp into `(−1, 1)` — the calibration target is a smile
        // quantity, not a correlation, so a slight boundary violation
        // is possible when the pillar grid is sparse. A hard clamp
        // with a logged warning is benign; callers that need strict
        // bounds should add a pillar.
        let rho_i_clamped = rho_i.clamp(-0.999, 0.999);
        rho_values.push(rho_i_clamped);
    }
    let rho_sched = PiecewiseConstant::new(knots, rho_values);

    let params = TimeDependentSabrParams::new(alpha_sched, rho_sched, nu_sched, beta, forward_0);

    // --- Post-fit diagnostics ---------------------------------------
    let per_pillar: Vec<PillarDiagnostics> = pillars
        .iter()
        .zip(&market)
        .zip(&stage1_rmse)
        .enumerate()
        .map(|(i, ((pi, mp), r1))| {
            let t = pi.expiry;
            let sub_nu = PiecewiseConstant::new(
                params.alpha.knots[..=i + 1].to_vec(),
                params.nu.values[..=i].to_vec(),
            );
            let sub_alpha = PiecewiseConstant::new(
                params.alpha.knots[..=i + 1].to_vec(),
                params.alpha.values[..=i].to_vec(),
            );
            let sub_rho = PiecewiseConstant::new(
                params.alpha.knots[..=i + 1].to_vec(),
                params.rho.values[..=i].to_vec(),
            );
            let mod_nu = effective_vol_vol(&sub_nu, &sub_alpha, t);
            let mod_omega = effective_term_structure(&sub_nu, &sub_alpha, t);
            let mod_rho = effective_correlation(&sub_nu, &sub_alpha, &sub_rho, t);
            let errs = [mod_nu - mp.nu, mod_omega - mp.alpha, mod_rho - mp.rho];
            let post_fit_rmse = (errs.iter().map(|e| e * e).sum::<f64>() / 3.0).sqrt();
            PillarDiagnostics {
                expiry: pi.expiry,
                market_effective: *mp,
                stage1_rmse: *r1,
                post_fit_rmse,
            }
        })
        .collect();

    TimeDependentCalibrationResult { params, per_pillar }
}

/// Bisect a monotone-increasing `f` on `[lo, hi]` for the root of
/// `f(x) = 0`. If the bracket isn't sign-reversed (target outside the
/// reachable range on this segment) returns the nearest endpoint —
/// calibration proceeds with a known residual rather than panicking.
/// The residual is captured downstream via `post_fit_rmse`.
fn bisect<F: Fn(f64) -> f64>(f: F, mut lo: f64, mut hi: f64, tol: f64, max_iter: usize) -> f64 {
    let f_lo = f(lo);
    let f_hi = f(hi);
    if f_lo > 0.0 {
        // Target below lo; best-effort = lo.
        return lo;
    }
    if f_hi < 0.0 {
        // Target above hi even after widening; return hi.
        return hi;
    }
    for _ in 0..max_iter {
        let mid = 0.5 * (lo + hi);
        let f_mid = f(mid);
        if f_mid.abs() < tol || (hi - lo) < tol {
            return mid;
        }
        if f_mid < 0.0 {
            lo = mid;
        } else {
            hi = mid;
        }
    }
    0.5 * (lo + hi)
}

/// Trait-object wrapper for the 4-stage time-dependent SABR
/// surface calibrator. `Market` is the multi-pillar
/// `Vec<MarketSmileStrip>` shape; the trait drives
/// [`calibrate_time_dependent`] after building `PillarTarget`s from
/// each strip.
///
/// `forward_0` is used as the initial forward in the resulting
/// [`TimeDependentSabrParams`]; downstream simulator builders
/// typically override this per-pillar via
/// [`TimeDependentSabrParams::new`] when rerunning MC at a specific
/// expiry.
pub struct SabrTimeDependentSurfaceCalibrator {
    pub beta: f64,
    pub forward_0: f64,
}

impl Calibration for SabrTimeDependentSurfaceCalibrator {
    type Market = Vec<MarketSmileStrip>;
    type Params = TimeDependentSabrParams;

    fn calibrate(
        &self,
        pillars: &Self::Market,
        options: NelderMeadOptions,
    ) -> Result<CalibrationReport<Self::Params>> {
        let pillar_targets: Vec<PillarTarget> = pillars.iter().map(|s| s.pillar_target()).collect();
        let res = calibrate_time_dependent(&pillar_targets, self.beta, self.forward_0, options);
        // Surface-level RMSE: max across pillars of the stage-1 smile
        // RMSE. This is the most conservative single summary number;
        // callers who need per-pillar breakdowns should call
        // `calibrate_time_dependent` directly and consume
        // `per_pillar`.
        let rmse = res
            .per_pillar
            .iter()
            .map(|d| d.stage1_rmse)
            .fold(0.0_f64, f64::max);
        Ok(CalibrationReport {
            params: res.params,
            rmse,
            // Multi-stage sequential calibrator: no single optimiser
            // diagnostic is representative.
            optimiser: None,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::forex::sabr::hagan_implied_vol;

    fn options() -> NelderMeadOptions {
        NelderMeadOptions {
            max_iter: 600,
            ftol: 1.0e-10,
            xtol: 1.0e-8,
            step_frac: 0.10,
        }
    }

    /// Fully-constant schedule round-trip: one pillar ⇒ stage-1 fit
    /// reproduces the truth, and stages 2-4 trivially give the same
    /// constant on the single segment.
    #[test]
    fn single_pillar_recovers_constant_sabr() {
        let truth = SabrParams::new(0.15, 0.5, -0.30, 0.45);
        let expiry = 1.0_f64;
        let f0 = 1.30_f64;
        let strikes = vec![1.05, 1.15, 1.30, 1.50, 1.70];
        let vols: Vec<f64> = strikes
            .iter()
            .map(|&k| hagan_implied_vol(&truth, f0, k, expiry))
            .collect();
        let pillars = vec![PillarTarget {
            expiry,
            forward: f0,
            strikes,
            market_vols: vols,
        }];
        let res = calibrate_time_dependent(&pillars, 0.5, f0, options());
        assert_eq!(res.params.alpha.values.len(), 1);
        let alpha = res.params.alpha.values[0];
        let rho = res.params.rho.values[0];
        let nu = res.params.nu.values[0];
        assert!(
            (alpha - truth.alpha).abs() < 1e-3,
            "α: {} vs {}",
            alpha,
            truth.alpha
        );
        assert!(
            (rho - truth.rho).abs() < 1e-3,
            "ρ: {} vs {}",
            rho,
            truth.rho
        );
        assert!((nu - truth.nu).abs() < 1e-3, "ν: {} vs {}", nu, truth.nu);
    }

    /// Two-pillar synthetic round-trip: the *actual* calibration
    /// success criterion is that the calibrated time-dependent model
    /// reprices the market smiles to within a small vol tolerance
    /// (the `post_fit_rmse` diagnostic tracks effective-parameter
    /// match, which the stage-2 approximation leaves as a
    /// second-order residual).
    ///
    /// We check that model Hagan IVs using the *model-effective*
    /// parameters at each pillar match the market smile vols within
    /// 50 bp RMSE — loose enough to absorb the stage-2 `ω(t) ≈
    /// ω̃ᵢ_mar` approximation (paper §5.2 reports ~30 bp errors from
    /// this approximation without the Phase-5 local-vol compensator)
    /// and tight enough to catch regressions in any of the three
    /// sequential sub-stages.
    #[test]
    fn two_pillar_round_trip_reprices_market_smiles() {
        use crate::models::forex::sabr_effective::{
            effective_correlation, effective_term_structure, effective_vol_vol,
        };
        let beta = 0.5_f64;
        let f0 = 1.30_f64;
        let knots = vec![0.0, 1.0, 2.0];
        let alpha_true = PiecewiseConstant::new(knots.clone(), vec![0.15, 0.12]);
        let rho_true = PiecewiseConstant::new(knots.clone(), vec![-0.30, -0.20]);
        let nu_true = PiecewiseConstant::new(knots.clone(), vec![0.45, 0.35]);

        let strikes = vec![1.05, 1.15, 1.30, 1.50, 1.70];
        let mut pillars = Vec::new();
        for &t in &[1.0_f64, 2.0] {
            let g = effective_vol_vol(&nu_true, &alpha_true, t);
            let w = effective_term_structure(&nu_true, &alpha_true, t);
            let r = effective_correlation(&nu_true, &alpha_true, &rho_true, t);
            let p = SabrParams::new(w, beta, r, g);
            let vols: Vec<f64> = strikes
                .iter()
                .map(|&k| hagan_implied_vol(&p, f0, k, t))
                .collect();
            pillars.push(PillarTarget {
                expiry: t,
                forward: f0,
                strikes: strikes.clone(),
                market_vols: vols,
            });
        }

        let res = calibrate_time_dependent(&pillars, beta, f0, options());

        // At each pillar, derive the model-effective SABR and price
        // market strikes with Hagan; check RMSE vs market vols.
        for (i, pi) in pillars.iter().enumerate() {
            let t = pi.expiry;
            let sub_nu = PiecewiseConstant::new(
                res.params.nu.knots[..=i + 1].to_vec(),
                res.params.nu.values[..=i].to_vec(),
            );
            let sub_alpha = PiecewiseConstant::new(
                res.params.alpha.knots[..=i + 1].to_vec(),
                res.params.alpha.values[..=i].to_vec(),
            );
            let sub_rho = PiecewiseConstant::new(
                res.params.rho.knots[..=i + 1].to_vec(),
                res.params.rho.values[..=i].to_vec(),
            );
            let mod_g = effective_vol_vol(&sub_nu, &sub_alpha, t);
            let mod_w = effective_term_structure(&sub_nu, &sub_alpha, t);
            let mod_r = effective_correlation(&sub_nu, &sub_alpha, &sub_rho, t);
            let eff = SabrParams::new(mod_w, beta, mod_r, mod_g);
            let mut ssr = 0.0;
            for (k, mv) in pi.strikes.iter().zip(&pi.market_vols) {
                let iv = hagan_implied_vol(&eff, pi.forward, *k, t);
                ssr += (iv - mv).powi(2);
            }
            let rmse = (ssr / pi.strikes.len() as f64).sqrt();
            assert!(
                rmse < 50.0e-4,
                "T={}: repriced market smile RMSE {:.2} bp vol > 50 bp",
                t,
                rmse * 10_000.0,
            );
        }
    }

    /// Three-pillar, declining ν schedule (typical FX smile
    /// structure: curvature decreases at longer expiries). The
    /// calibrator should produce a ν schedule with non-increasing
    /// segments on synthetic input.
    #[test]
    fn declining_nu_schedule_is_recovered_monotone() {
        let beta = 0.5_f64;
        let f0 = 1.30_f64;
        let knots = vec![0.0, 1.0, 2.0, 3.0];
        let alpha_true = PiecewiseConstant::new(knots.clone(), vec![0.15; 3]);
        let rho_true = PiecewiseConstant::new(knots.clone(), vec![-0.30; 3]);
        let nu_true = PiecewiseConstant::new(knots.clone(), vec![0.60, 0.45, 0.30]);

        use crate::models::forex::sabr_effective::{
            effective_correlation, effective_term_structure, effective_vol_vol,
        };
        let strikes = vec![1.05, 1.20, 1.30, 1.45, 1.70];
        let mut pillars = Vec::new();
        for &t in &[1.0_f64, 2.0, 3.0] {
            let g = effective_vol_vol(&nu_true, &alpha_true, t);
            let w = effective_term_structure(&nu_true, &alpha_true, t);
            let r = effective_correlation(&nu_true, &alpha_true, &rho_true, t);
            let p = SabrParams::new(w, beta, r, g);
            let vols: Vec<f64> = strikes
                .iter()
                .map(|&k| hagan_implied_vol(&p, f0, k, t))
                .collect();
            pillars.push(PillarTarget {
                expiry: t,
                forward: f0,
                strikes: strikes.clone(),
                market_vols: vols,
            });
        }
        let res = calibrate_time_dependent(&pillars, beta, f0, options());
        let nu_fit = &res.params.nu.values;
        // Monotone non-increasing within a small tolerance.
        for w in nu_fit.windows(2) {
            assert!(w[1] <= w[0] + 1e-3, "ν not monotone: {} > {}", w[1], w[0]);
        }
    }

    /// Pillars out of order panic — API contract.
    #[test]
    #[should_panic(expected = "sorted")]
    fn unsorted_pillars_panic() {
        let pillars = vec![
            PillarTarget {
                expiry: 2.0,
                forward: 1.0,
                strikes: vec![1.0],
                market_vols: vec![0.1],
            },
            PillarTarget {
                expiry: 1.0,
                forward: 1.0,
                strikes: vec![1.0],
                market_vols: vec![0.1],
            },
        ];
        let _ = calibrate_time_dependent(&pillars, 0.5, 1.0, options());
    }
}