quantsupport 0.1.7

use crate::{
    ad::{constant::Const, dual::DualFwd, expr::FloatExt, scalar::Scalar},
    core::{elements::curveelement::ADCurveElement, pillars::Pillars},
    math::interpolation::interpolator::{Interpolate, Interpolator},
    rates::{
        compounding::Compounding, interestrate::InterestRate,
        yieldtermstructure::interestratestermstructure::InterestRatesTermStructure,
    },
    time::{date::Date, daycounter::DayCounter, enums::Frequency},
    utils::errors::{QSError, Result},
};

/// A discount factors term structure.
///
/// ## Example
/// ```rust
/// use quantsupport::prelude::*;
///
/// let dates = vec![
///     Date::new(2020, 1, 1),
///     Date::new(2020, 4, 1),
///     Date::new(2020, 7, 1),
///     Date::new(2020, 10, 1),
///     Date::new(2021, 1, 1),
/// ];
///
/// let discount_factors = vec![1.0, 0.99, 0.98, 0.97, 0.96];
/// let day_counter = DayCounter::Actual360;
///
/// let discount_term_structure = DiscountTermStructure::<f64>::new(
///    dates.clone(),
///    discount_factors.clone(),
///    day_counter,
///    Interpolator::Linear,
///    true).unwrap();
///
/// assert_eq!(
///     discount_term_structure.dates().clone(),
///     dates
/// );
/// assert_eq!(
///     discount_term_structure.discount_factors().clone(),
///     discount_factors
/// );
///  ```
/// A discount factors term structure.
#[derive(Clone)]
pub struct DiscountTermStructure<T>
where
    T: Scalar,
{
    reference_date: Date,
    dates: Vec<Date>,
    year_fractions: Vec<f64>,
    discount_factors: Vec<T>,
    interpolator: Interpolator,
    day_counter: DayCounter,
    enable_extrapolation: bool,
    pillar_labels: Option<Vec<String>>,
    pillar_values: Option<Vec<T>>,
    ift_sensitivities: Option<Vec<Vec<f64>>>,
}

impl<T> DiscountTermStructure<T>
where
    T: Scalar,
{
    /// Returns a reference to the vector of dates.
    #[must_use]
    pub const fn dates(&self) -> &Vec<Date> {
        &self.dates
    }

    /// Returns a reference to the vector of discount factors.
    #[must_use]
    pub const fn discount_factors(&self) -> &Vec<T> {
        &self.discount_factors
    }

    /// Returns the day counter convention used.
    #[must_use]
    pub const fn day_counter(&self) -> DayCounter {
        self.day_counter
    }

    /// Returns whether extrapolation is enabled.
    #[must_use]
    pub const fn enable_extrapolation(&self) -> bool {
        self.enable_extrapolation
    }

    /// Returns the interpolator used.
    #[must_use]
    pub const fn interpolator(&self) -> Interpolator {
        self.interpolator
    }

    /// Returns the reference date of the term structure.
    ///
    /// ## Errors
    /// Returns an error if the reference date is not set.
    pub fn with_pillar_labels(mut self, pillar_labels: Vec<String>) -> Result<Self> {
        let expected_len = self
            .pillar_values
            .as_ref()
            .map_or(self.discount_factors.len(), Vec::len);
        if pillar_labels.len() != expected_len {
            return Err(QSError::InvalidValueErr(
                "Pillar labels need to have the same size as the exposed pillar values".into(),
            ));
        }
        self.pillar_labels = Some(pillar_labels);
        Ok(self)
    }

    /// Stores the IFT sensitivity matrix so that `put_pillars_on_tape` can
    /// rebuild AD-connected discount factors.
    #[must_use]
    pub fn with_ift_sensitivities(mut self, sensitivities: Vec<Vec<f64>>) -> Self {
        self.ift_sensitivities = Some(sensitivities);
        self
    }

    /// Overrides the values exposed through the [`Pillars`] interface.
    ///
    /// # Errors
    /// Returns an error if the provided values do not match the number of curve nodes.
    pub fn with_pillar_values(mut self, pillar_values: Vec<T>) -> Result<Self> {
        if let Some(pillar_labels) = &self.pillar_labels {
            if pillar_values.len() != pillar_labels.len() {
                return Err(QSError::InvalidValueErr(
                    "Pillar values need to have the same size as pillar labels".into(),
                ));
            }
        }
        if pillar_values.is_empty() {
            return Err(QSError::InvalidValueErr(
                "Pillar values cannot be empty".into(),
            ));
        }
        self.pillar_values = Some(pillar_values);
        Ok(self)
    }
}

impl DiscountTermStructure<f64> {
    /// Creates a new `DiscountTermStructure` with the given dates, discount factors, day counter, interpolator, and extrapolation setting.
    ///
    /// ## Arguments
    ///
    /// * `dates` - Vector of dates for the discount factors
    /// * `discount_factors` - Vector of discount factors corresponding to the dates
    /// * `day_counter` - Day counter convention to use for year fraction calculations
    /// * `interpolator` - Interpolation method to use
    /// * `enable_extrapolation` - Whether to allow extrapolation beyond the given dates
    ///
    /// ## Errors
    ///
    /// Returns an error if dates and discount factors have different lengths or if the first discount factor is not 1.0.
    pub fn new(
        dates: Vec<Date>,
        discount_factors: Vec<f64>,
        day_counter: DayCounter,
        interpolator: Interpolator,
        enable_extrapolation: bool,
    ) -> Result<Self> {
        // check if year_fractions and discount_factors have the same size
        if dates.len() != discount_factors.len() {
            return Err(QSError::InvalidValueErr(
                "Dates and discount_factors need to have the same size".to_string(),
            ));
        }

        // order dates y discount_factors
        let mut zipped = dates.into_iter().zip(discount_factors).collect::<Vec<_>>();
        zipped.sort_by_key(|a| a.0);
        let (dates, discount_factors): (Vec<Date>, Vec<f64>) = zipped.into_iter().unzip();

        // discount_factors[0] needs to be 1.0
        if (discount_factors[0] - 1.0).abs() > 1e-12 {
            return Err(QSError::InvalidValueErr(
                "First discount factor needs to be 1.0".to_string(),
            ));
        }
        let reference_date = dates[0];
        let year_fractions: Vec<f64> = dates
            .iter()
            .map(|x| day_counter.year_fraction(reference_date, *x))
            .collect();

        Ok(Self {
            reference_date,
            dates,
            year_fractions,
            discount_factors,
            interpolator,
            day_counter,
            enable_extrapolation,
            pillar_labels: None,
            pillar_values: None,
            ift_sensitivities: None,
        })
    }
}

impl DiscountTermStructure<DualFwd> {
    /// Creates a new `DiscountTermStructure` with the given dates, discount factors, day counter, interpolator, and extrapolation setting.
    ///
    /// # Arguments
    ///
    /// * `dates` - Vector of dates for the discount factors
    /// * `discount_factors` - Vector of discount factors corresponding to the dates
    /// * `day_counter` - Day counter convention to use for year fraction calculations
    /// * `interpolator` - Interpolation method to use
    /// * `enable_extrapolation` - Whether to allow extrapolation beyond the given dates
    ///
    /// # Errors
    ///
    /// Returns an error if dates and discount factors have different lengths or if the first discount factor is not 1.0.
    pub fn new(
        dates: Vec<Date>,
        discount_factors: Vec<DualFwd>,
        day_counter: DayCounter,
        interpolator: Interpolator,
        enable_extrapolation: bool,
    ) -> Result<Self> {
        // check if year_fractions and discount_factors have the same size
        if dates.len() != discount_factors.len() {
            return Err(QSError::InvalidValueErr(
                "Dates and discount_factors need to have the same size".to_string(),
            ));
        }

        // order dates y discount_factors
        let mut zipped = dates.into_iter().zip(discount_factors).collect::<Vec<_>>();
        zipped.sort_by_key(|a| a.0);
        let (dates, discount_factors): (Vec<Date>, Vec<DualFwd>) = zipped.into_iter().unzip();

        // discount_factors[0] needs to be 1.0
        if (discount_factors[0] - Const::one()).abs() > 1e-12 {
            return Err(QSError::InvalidValueErr(
                "First discount factor needs to be 1.0".to_string(),
            ));
        }
        let reference_date = dates[0];
        let year_fractions: Vec<f64> = dates
            .iter()
            .map(|x| day_counter.year_fraction(reference_date, *x))
            .collect();

        Ok(Self {
            reference_date,
            dates,
            year_fractions,
            discount_factors,
            interpolator,
            day_counter,
            enable_extrapolation,
            pillar_labels: None,
            pillar_values: None,
            ift_sensitivities: None,
        })
    }
}

impl Pillars<DualFwd> for DiscountTermStructure<DualFwd> {
    fn pillars(&self) -> Option<Vec<(String, &DualFwd)>> {
        self.pillar_labels.as_ref().map(|labels| {
            let values = self
                .pillar_values
                .as_ref()
                .unwrap_or(&self.discount_factors);
            labels.iter().cloned().zip(values.iter()).collect()
        })
    }

    fn pillar_labels(&self) -> Option<Vec<String>> {
        self.pillar_labels.clone()
    }

    fn put_pillars_on_tape(&mut self) {
        if let Some(pillar_values) = self.pillar_values.as_mut() {
            pillar_values.iter_mut().for_each(DualFwd::put_on_tape);

            // When IFT sensitivities are stored, rebuild discount_factors as
            // tape-connected expressions:
            //   DF(i+1) = df_val + Σ_j S[i][j] * (q_j - q_j_val)
            // The numeric result is exact (delta evaluates to 0) but the AD
            // graph now connects discount_factors → pillar_values.
            // The matrix may be non-square when cross-curve dependencies
            // have been pre-composed at bootstrap time (n_pillars > n_dfs).
            if let Some(ref sens) = self.ift_sensitivities {
                let n_pillars = pillar_values.len();
                let mut new_dfs = Vec::with_capacity(sens.len() + 1);
                new_dfs.push(DualFwd::new(1.0)); // DF(0) = 1
                for (i, sens_row) in sens.iter().enumerate() {
                    let mut df_ad = DualFwd::new(self.discount_factors[i + 1].value());
                    for j in 0..n_pillars {
                        let s = sens_row[j];
                        if s.abs() > 1e-16 {
                            let delta: DualFwd =
                                (pillar_values[j] - DualFwd::new(pillar_values[j].value())).into();
                            df_ad = (df_ad + DualFwd::new(s) * delta).into();
                        }
                    }
                    new_dfs.push(df_ad);
                }
                self.discount_factors = new_dfs;
            }
        } else {
            self.discount_factors
                .iter_mut()
                .for_each(DualFwd::put_on_tape);
        }
    }
}

impl InterestRatesTermStructure<f64> for DiscountTermStructure<f64> {
    fn reference_date(&self) -> Date {
        self.reference_date
    }

    fn discount_factor(&self, date: Date) -> Result<f64> {
        if date < self.reference_date() {
            return Err(QSError::InvalidValueErr(
                "Date needs to be greater than reference date".to_string(),
            ));
        }
        if date == self.reference_date() {
            return Ok(1.0);
        }

        let year_fraction = self
            .day_counter()
            .year_fraction(self.reference_date(), date);

        let discount_factor = self.interpolator.interpolate(
            year_fraction,
            &self.year_fractions,
            &self.discount_factors,
            self.enable_extrapolation,
        )?;
        Ok(discount_factor)
    }

    fn forward_rate(
        &self,
        start_date: Date,
        end_date: Date,
        comp: Compounding,
        freq: Frequency,
    ) -> Result<f64> {
        let discount_factor_to_star = self.discount_factor(start_date)?;
        let discount_factor_to_end = self.discount_factor(end_date)?;

        let comp_factor = discount_factor_to_star / discount_factor_to_end;
        let t = self.day_counter().year_fraction(start_date, end_date);

        Ok(
            InterestRate::<f64>::implied_rate(comp_factor, self.day_counter(), comp, freq, t)?
                .rate(),
        )
    }

    fn nodes(&self) -> Option<Vec<(Date, f64)>> {
        Some(
            self.dates
                .iter()
                .copied()
                .zip(self.discount_factors.iter().copied())
                .collect(),
        )
    }

    fn day_counter(&self) -> Option<DayCounter> {
        Some(self.day_counter)
    }

    fn discount_factor_from_time(&self, t: f64) -> Result<f64> {
        if t < 0.0 {
            return Err(QSError::InvalidValueErr(
                "Time must be non-negative".to_string(),
            ));
        }
        if t == 0.0 {
            return Ok(1.0);
        }
        self.interpolator.interpolate(
            t,
            &self.year_fractions,
            &self.discount_factors,
            self.enable_extrapolation,
        )
    }

    fn forward_rate_from_time(&self, start: f64, end: f64) -> Result<f64> {
        let dt = end - start;
        let (s, e, dt) = if dt.abs() < 1e-10 {
            let eps = 1e-6;
            (start, start + eps, eps)
        } else {
            (start, end, dt)
        };
        let df_start = self.discount_factor_from_time(s)?;
        let df_end = self.discount_factor_from_time(e)?;
        let fwd = (df_start / df_end - 1.0) / dt;
        Ok(fwd)
    }
}

impl InterestRatesTermStructure<DualFwd> for DiscountTermStructure<DualFwd> {
    fn reference_date(&self) -> Date {
        self.reference_date
    }
    fn discount_factor(&self, date: Date) -> Result<DualFwd> {
        if date < self.reference_date() {
            return Err(QSError::InvalidValueErr(
                "Date needs to be greater than reference date".to_string(),
            ));
        }
        if date == self.reference_date() {
            return Ok(DualFwd::one());
        }

        let year_fraction = DualFwd::new(
            self.day_counter()
                .year_fraction(self.reference_date(), date),
        );

        let tmp_yfs = self
            .year_fractions
            .iter()
            .copied()
            .map(DualFwd::new)
            .collect::<Vec<DualFwd>>();

        let discount_factor = self.interpolator.interpolate(
            year_fraction,
            &tmp_yfs,
            &self.discount_factors,
            self.enable_extrapolation,
        )?;
        Ok(discount_factor)
    }

    fn forward_rate(
        &self,
        start_date: Date,
        end_date: Date,
        comp: Compounding,
        freq: Frequency,
    ) -> Result<DualFwd> {
        let discount_factor_to_star = self.discount_factor(start_date)?;
        let discount_factor_to_end = self.discount_factor(end_date)?;

        let comp_factor = discount_factor_to_star / discount_factor_to_end;
        let t = self.day_counter().year_fraction(start_date, end_date);

        Ok(InterestRate::<DualFwd>::implied_rate(
            comp_factor.into(),
            self.day_counter(),
            comp,
            freq,
            t,
        )?
        .rate())
    }

    fn nodes(&self) -> Option<Vec<(Date, DualFwd)>> {
        Some(
            self.dates
                .iter()
                .copied()
                .zip(self.discount_factors.iter().copied())
                .collect(),
        )
    }

    fn day_counter(&self) -> Option<DayCounter> {
        Some(self.day_counter)
    }

    fn discount_factor_from_time(&self, t: f64) -> Result<DualFwd> {
        if t < 0.0 {
            return Err(QSError::InvalidValueErr(
                "Time must be non-negative".to_string(),
            ));
        }
        if t == 0.0 {
            return Ok(DualFwd::one());
        }
        let yf = DualFwd::new(t);
        let tmp_yfs: Vec<DualFwd> = self
            .year_fractions
            .iter()
            .copied()
            .map(DualFwd::new)
            .collect();
        self.interpolator.interpolate(
            yf,
            &tmp_yfs,
            &self.discount_factors,
            self.enable_extrapolation,
        )
    }

    fn forward_rate_from_time(&self, start: f64, end: f64) -> Result<DualFwd> {
        let dt = end - start;
        let (s, e, dt) = if dt.abs() < 1e-10 {
            let eps = 1e-6;
            (start, start + eps, eps)
        } else {
            (start, end, dt)
        };
        let df_start = self.discount_factor_from_time(s)?;
        let df_end = self.discount_factor_from_time(e)?;
        let fwd = (df_start / df_end - 1.0) / dt;
        Ok(fwd.into())
    }
}

impl ADCurveElement for DiscountTermStructure<DualFwd> {
    fn ift_sensitivities(&self) -> Option<&[Vec<f64>]> {
        self.ift_sensitivities.as_deref()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_year_fractions() {
        let dates = vec![
            Date::new(2020, 1, 1),
            Date::new(2020, 4, 1),
            Date::new(2020, 7, 1),
            Date::new(2020, 10, 1),
            Date::new(2021, 1, 1),
        ];
        let discount_factors = vec![1.0, 0.99, 0.98, 0.97, 0.96];
        let day_counter = DayCounter::Actual360;

        let discount_term_structure = DiscountTermStructure::<f64>::new(
            dates,
            discount_factors,
            day_counter,
            Interpolator::Linear,
            true,
        )
        .unwrap_or_else(|e| {
            panic!("DiscountTermStructure::new should succeed in test_year_fractions: {e}")
        });

        assert_eq!(
            discount_term_structure.dates(),
            &vec![
                Date::new(2020, 1, 1),
                Date::new(2020, 4, 1),
                Date::new(2020, 7, 1),
                Date::new(2020, 10, 1),
                Date::new(2021, 1, 1)
            ]
        );
    }

    #[test]
    fn test_discount_dactors() {
        let dates = vec![
            Date::new(2020, 1, 1),
            Date::new(2020, 4, 1),
            Date::new(2020, 7, 1),
            Date::new(2020, 10, 1),
            Date::new(2021, 1, 1),
        ];
        let discount_factors = vec![1.0, 0.99, 0.98, 0.97, 0.96];
        let day_counter = DayCounter::Actual360;

        let discount_term_structure = DiscountTermStructure::<f64>::new(
            dates,
            discount_factors,
            day_counter,
            Interpolator::Linear,
            true,
        )
        .unwrap_or_else(|e| {
            panic!("DiscountTermStructure::new should succeed in test_discount_dactors: {e}")
        });

        assert_eq!(
            discount_term_structure.discount_factors(),
            &vec![1.0, 0.99, 0.98, 0.97, 0.96]
        );
    }

    #[test]
    fn test_reference_date() {
        let dates = vec![
            Date::new(2020, 1, 1),
            Date::new(2020, 4, 1),
            Date::new(2020, 7, 1),
            Date::new(2020, 10, 1),
            Date::new(2021, 1, 1),
        ];
        let discount_factors = vec![1.0, 0.99, 0.98, 0.97, 0.96];
        let day_counter = DayCounter::Actual360;

        let discount_term_structure = DiscountTermStructure::<f64>::new(
            dates,
            discount_factors,
            day_counter,
            Interpolator::Linear,
            true,
        )
        .unwrap_or_else(|e| {
            panic!("DiscountTermStructure::new should succeed in test_reference_date: {e}")
        });

        assert_eq!(
            discount_term_structure.reference_date(),
            Date::new(2020, 1, 1)
        );
    }

    #[test]
    fn test_interpolation() {
        let dates = vec![
            Date::new(2020, 1, 1),
            Date::new(2020, 4, 1),
            Date::new(2020, 7, 1),
            Date::new(2020, 10, 1),
            Date::new(2021, 1, 1),
        ];
        let discount_factors = vec![1.0, 0.99, 0.98, 0.97, 0.96];
        let day_counter = DayCounter::Actual360;

        let discount_term_structure = DiscountTermStructure::<f64>::new(
            dates,
            discount_factors,
            day_counter,
            Interpolator::Linear,
            true,
        )
        .unwrap_or_else(|e| {
            panic!("DiscountTermStructure::new should succeed in test_interpolation: {e}")
        });

        let df = discount_term_structure
            .discount_factor(Date::new(2020, 6, 1))
            .unwrap_or_else(|e| panic!("discount_factor failed in test_interpolation: {e}"));
        assert!((df - 0.9832967032967033).abs() < 1e-8);
    }

    #[test]

    fn test_forward_rate() {
        let dates = vec![
            Date::new(2020, 1, 1),
            Date::new(2020, 4, 1),
            Date::new(2020, 7, 1),
            Date::new(2020, 10, 1),
            Date::new(2021, 1, 1),
        ];
        let discount_factors = vec![1.0, 0.99, 0.98, 0.97, 0.96];
        let day_counter = DayCounter::Actual360;

        let discount_term_structure = DiscountTermStructure::<f64>::new(
            dates,
            discount_factors,
            day_counter,
            Interpolator::Linear,
            true,
        )
        .unwrap_or_else(|e| {
            panic!("DiscountTermStructure::new should succeed in test_forward_rate: {e}")
        });

        let comp = Compounding::Simple;
        let freq = Frequency::Annual;

        let fwd = discount_term_structure
            .forward_rate(Date::new(2020, 1, 1), Date::new(2020, 12, 31), comp, freq)
            .unwrap_or_else(|e| panic!("forward_rate failed in test_forward_rate: {e}"));
        assert!((fwd - 0.04097957689796514).abs() < 1e-8);
        println!("forward_rate: {fwd}");
    }

    #[test]
    fn order_dates() {
        let dates = vec![
            Date::new(2020, 1, 1),
            Date::new(2020, 4, 1),
            Date::new(2020, 7, 1),
            Date::new(2020, 10, 1),
            Date::new(2021, 1, 1),
        ];
        let discount_factors = vec![0.99, 0.98, 0.97, 0.96, 1.0];
        let day_counter = DayCounter::Actual360;

        let discount_term_structure = DiscountTermStructure::<f64>::new(
            dates,
            discount_factors,
            day_counter,
            Interpolator::Linear,
            true,
        );

        assert!(discount_term_structure.is_err());
    }
}