stochastic-rs-quant 2.2.0

//! Schedule generation for coupon and payment dates.
//!
//! Generates periodic date schedules with business day adjustment,
//! stub handling, and end-of-month preservation.
//!
//! Reference: ISDA 2006 Definitions, Sections 4.15–4.16.

use chrono::Datelike;
use chrono::NaiveDate;

use super::business_day::BusinessDayConvention;
use super::day_count::DayCountConvention;
use super::day_count::days_in_month;
use super::holiday::Calendar;
use crate::traits::FloatExt;

/// Payment / coupon frequency.
#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Frequency {
  Annual,
  #[default]
  SemiAnnual,
  Quarterly,
  Monthly,
}

impl std::fmt::Display for Frequency {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    match self {
      Self::Annual => write!(f, "Annual"),
      Self::SemiAnnual => write!(f, "Semi-Annual"),
      Self::Quarterly => write!(f, "Quarterly"),
      Self::Monthly => write!(f, "Monthly"),
    }
  }
}

impl Frequency {
  /// Number of months per period.
  pub fn months(self) -> i32 {
    match self {
      Self::Annual => 12,
      Self::SemiAnnual => 6,
      Self::Quarterly => 3,
      Self::Monthly => 1,
    }
  }

  /// Number of coupon periods per year.
  pub fn periods_per_year(self) -> u32 {
    match self {
      Self::Annual => 1,
      Self::SemiAnnual => 2,
      Self::Quarterly => 4,
      Self::Monthly => 12,
    }
  }
}

/// Direction of date generation.
#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DateGenerationRule {
  /// Generate dates from the effective date toward the termination date.
  Forward,
  /// Generate dates from the termination date toward the effective date.
  #[default]
  Backward,
}

impl std::fmt::Display for DateGenerationRule {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    match self {
      Self::Forward => write!(f, "Forward"),
      Self::Backward => write!(f, "Backward"),
    }
  }
}

/// A generated schedule of dates.
#[derive(Debug, Clone)]
pub struct Schedule {
  /// Unadjusted schedule dates.
  pub dates: Vec<NaiveDate>,
  /// Business-day-adjusted schedule dates.
  pub adjusted_dates: Vec<NaiveDate>,
}

impl Schedule {
  /// Compute year fractions between consecutive adjusted dates.
  pub fn year_fractions<T: FloatExt>(&self, convention: DayCountConvention) -> Vec<T> {
    self
      .adjusted_dates
      .windows(2)
      .map(|w| convention.year_fraction(w[0], w[1]))
      .collect()
  }
}

/// Stub period convention for schedules whose total length is not an
/// integer multiple of the payment period. Per ISDA 2006 §4.15.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StubConvention {
  /// Short stub at the start of the schedule (the first period is shorter
  /// than the regular period). Default for [`DateGenerationRule::Backward`].
  ShortFirst,
  /// Long stub at the start of the schedule (the first period absorbs both
  /// the irregular remainder and one regular period).
  LongFirst,
  /// Short stub at the end of the schedule. Default for
  /// [`DateGenerationRule::Forward`].
  ShortLast,
  /// Long stub at the end of the schedule.
  LongLast,
}

/// Fluent builder for [`Schedule`].
#[derive(Debug, Clone)]
pub struct ScheduleBuilder {
  effective: NaiveDate,
  termination: NaiveDate,
  frequency: Frequency,
  calendar: Option<Calendar>,
  convention: BusinessDayConvention,
  rule: DateGenerationRule,
  end_of_month: bool,
  /// Explicit stub convention. When `None`, defaults to `ShortFirst` for
  /// backward generation and `ShortLast` for forward generation (the
  /// implicit pre-rc.2 behaviour).
  stub: Option<StubConvention>,
}

impl ScheduleBuilder {
  pub fn new(effective: NaiveDate, termination: NaiveDate) -> Self {
    Self {
      effective,
      termination,
      frequency: Frequency::SemiAnnual,
      calendar: None,
      convention: BusinessDayConvention::ModifiedFollowing,
      rule: DateGenerationRule::Backward,
      end_of_month: false,
      stub: None,
    }
  }

  /// Set the stub convention. By default (no call), the builder picks
  /// `ShortFirst` for backward generation and `ShortLast` for forward
  /// generation; pass `LongFirst` / `LongLast` to merge the irregular
  /// stub with the next regular period instead.
  pub fn stub(mut self, stub: StubConvention) -> Self {
    self.stub = Some(stub);
    self
  }

  pub fn frequency(mut self, frequency: Frequency) -> Self {
    self.frequency = frequency;
    self
  }

  pub fn calendar(mut self, calendar: Calendar) -> Self {
    self.calendar = Some(calendar);
    self
  }

  pub fn convention(mut self, convention: BusinessDayConvention) -> Self {
    self.convention = convention;
    self
  }

  pub fn forward(mut self) -> Self {
    self.rule = DateGenerationRule::Forward;
    self
  }

  pub fn backward(mut self) -> Self {
    self.rule = DateGenerationRule::Backward;
    self
  }

  pub fn end_of_month(mut self, flag: bool) -> Self {
    self.end_of_month = flag;
    self
  }

  /// Build the schedule.
  pub fn build(self) -> Schedule {
    let period = self.frequency.months();
    let mut raw_dates = match self.rule {
      DateGenerationRule::Backward => {
        generate_backward(self.effective, self.termination, period, self.end_of_month)
      }
      DateGenerationRule::Forward => {
        generate_forward(self.effective, self.termination, period, self.end_of_month)
      }
    };

    raw_dates.sort();
    raw_dates.dedup();

    // Apply long-stub merging when the user has explicitly requested it.
    // Short-stub conventions (`ShortFirst` / `ShortLast`) match the default
    // `generate_*` output, so no extra work is needed in those cases.
    let stub = self.stub.unwrap_or(match self.rule {
      DateGenerationRule::Backward => StubConvention::ShortFirst,
      DateGenerationRule::Forward => StubConvention::ShortLast,
    });
    if raw_dates.len() >= 3 {
      match stub {
        StubConvention::LongFirst => {
          // Merge first stub with the next regular period: drop dates[1].
          // Only applies when there *is* a stub (start ≠ first regular date).
          let stub_period_months = months_between(raw_dates[0], raw_dates[1]);
          if stub_period_months != period {
            raw_dates.remove(1);
          }
        }
        StubConvention::LongLast => {
          // Merge last stub with the previous regular period: drop dates[n-2].
          let n = raw_dates.len();
          let stub_period_months = months_between(raw_dates[n - 2], raw_dates[n - 1]);
          if stub_period_months != period {
            raw_dates.remove(n - 2);
          }
        }
        StubConvention::ShortFirst | StubConvention::ShortLast => {}
      }
    }

    let adjusted = match &self.calendar {
      Some(cal) => raw_dates
        .iter()
        .map(|&d| self.convention.adjust(d, cal))
        .collect(),
      None => raw_dates.clone(),
    };

    Schedule {
      dates: raw_dates,
      adjusted_dates: adjusted,
    }
  }
}

/// Approximate calendar months between two dates (sign-preserving). Used
/// only to detect "is this a stub period?" — exact day-count is irrelevant.
fn months_between(a: NaiveDate, b: NaiveDate) -> i32 {
  use chrono::Datelike;
  let years = b.year() - a.year();
  let months = b.month() as i32 - a.month() as i32;
  years * 12 + months
}

fn generate_backward(
  effective: NaiveDate,
  termination: NaiveDate,
  period_months: i32,
  eom: bool,
) -> Vec<NaiveDate> {
  let mut dates = vec![termination];
  let mut i = 1i32;
  loop {
    let d = add_months(termination, -period_months * i, eom);
    if d <= effective {
      break;
    }
    dates.push(d);
    i += 1;
  }
  dates.push(effective);
  dates
}

fn generate_forward(
  effective: NaiveDate,
  termination: NaiveDate,
  period_months: i32,
  eom: bool,
) -> Vec<NaiveDate> {
  let mut dates = vec![effective];
  let mut i = 1i32;
  loop {
    let d = add_months(effective, period_months * i, eom);
    if d >= termination {
      break;
    }
    dates.push(d);
    i += 1;
  }
  dates.push(termination);
  dates
}

/// Add `months` calendar months to `date`, clamping to month-end if needed.
/// When `eom` is true and the input is the last day of its month, the result
/// is also the last day of the target month.
///
/// Uses the canonical [`days_in_month`](crate::calendar::day_count::days_in_month)
/// helper.
pub(crate) fn add_months(date: NaiveDate, months: i32, eom: bool) -> NaiveDate {
  let total = date.year() * 12 + date.month0() as i32 + months;
  let target_year = total.div_euclid(12);
  let target_month = (total.rem_euclid(12) + 1) as u32;
  let max_day = days_in_month(target_year, target_month);

  let day = if eom && date.day() == days_in_month(date.year(), date.month()) {
    max_day
  } else {
    date.day().min(max_day)
  };

  NaiveDate::from_ymd_opt(target_year, target_month, day).unwrap()
}

#[cfg(test)]
mod tests {
  use chrono::NaiveDate;

  use super::*;

  #[test]
  fn semiannual_two_year_schedule() {
    let s = ScheduleBuilder::new(
      NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
      NaiveDate::from_ymd_opt(2026, 1, 1).unwrap(),
    )
    .frequency(Frequency::SemiAnnual)
    .build();
    // 2 years semi-annual = 5 dates: t=0, +6m, +1y, +18m, +2y
    assert_eq!(s.dates.len(), 5);
  }

  #[test]
  fn frequency_periods_per_year() {
    assert_eq!(Frequency::Annual.periods_per_year(), 1);
    assert_eq!(Frequency::SemiAnnual.periods_per_year(), 2);
    assert_eq!(Frequency::Quarterly.periods_per_year(), 4);
    assert_eq!(Frequency::Monthly.periods_per_year(), 12);
  }

  #[test]
  fn frequency_months() {
    assert_eq!(Frequency::Annual.months(), 12);
    assert_eq!(Frequency::SemiAnnual.months(), 6);
    assert_eq!(Frequency::Quarterly.months(), 3);
  }
}