stochastic-rs-quant 2.6.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::NaiveDate;

use super::business_day::BusinessDayConvention;
use super::date_math::add_months;
use super::date_math::snap_to_imm;
use super::day_count::DayCountConvention;
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>,
  /// Snap every generated date to the nearest IMM date (3rd Wednesday of
  /// the same calendar quarter), per CME / LIFFE futures convention.
  imm: bool,
}

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,
      imm: false,
    }
  }

  /// Snap every generated date to its quarterly IMM date (3rd Wednesday
  /// of the enclosing March / June / September / December bucket). Used
  /// for futures-aligned schedules.
  pub fn imm(mut self, flag: bool) -> Self {
    self.imm = flag;
    self
  }

  /// Set the stub convention.
  ///
  /// **Defaults (when this method is not called):**
  /// - [`DateGenerationRule::Backward`] → [`StubConvention::ShortFirst`]
  /// - [`DateGenerationRule::Forward`] → [`StubConvention::ShortLast`]
  ///
  /// **Edge cases:**
  /// - `LongFirst` / `LongLast` are **no-ops** when the schedule has no stub
  ///   (i.e. when the total tenor is an integer multiple of the regular
  ///   period). Detection is by comparing `months_between(stub_endpoints)`
  ///   against the regular `period` — equal means no stub to merge.
  /// - Long-merge applies only when the raw schedule has ≥ 3 dates; a
  ///   2-date schedule (single period) is left untouched regardless of the
  ///   chosen convention.
  /// - Mixing a Long convention with the **opposite** direction (e.g.
  ///   `Forward` + `LongFirst`) is allowed and merges the start-side stub
  ///   if any was produced — useful when generating an end-anchored
  ///   schedule but wanting the *first* irregular period absorbed.
  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)
      }
    };

    if self.imm {
      for d in raw_dates.iter_mut() {
        *d = snap_to_imm(*d);
      }
    }

    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
}

#[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);
  }

  #[test]
  fn imm_quarterly_schedule_lands_on_third_wednesdays() {
    let s = ScheduleBuilder::new(
      NaiveDate::from_ymd_opt(2024, 3, 1).unwrap(),
      NaiveDate::from_ymd_opt(2025, 3, 1).unwrap(),
    )
    .frequency(Frequency::Quarterly)
    .imm(true)
    .build();
    use chrono::Datelike as _;
    for d in &s.dates {
      assert_eq!(d.weekday(), chrono::Weekday::Wed, "{d} not Wednesday");
      assert!((15..=21).contains(&d.day()), "{d} not in 3rd week");
      assert!(matches!(d.month(), 3 | 6 | 9 | 12), "{d} not in IMM month");
    }
  }

  #[test]
  fn imm_flag_disabled_keeps_generated_dates() {
    let s_plain = ScheduleBuilder::new(
      NaiveDate::from_ymd_opt(2024, 3, 1).unwrap(),
      NaiveDate::from_ymd_opt(2025, 3, 1).unwrap(),
    )
    .frequency(Frequency::Quarterly)
    .build();
    let s_imm = ScheduleBuilder::new(
      NaiveDate::from_ymd_opt(2024, 3, 1).unwrap(),
      NaiveDate::from_ymd_opt(2025, 3, 1).unwrap(),
    )
    .frequency(Frequency::Quarterly)
    .imm(true)
    .build();
    use chrono::Datelike as _;
    // Same length and final endpoint should match the IMM bucket of the
    // termination, but the generic schedule keeps day-of-month 1.
    assert_eq!(s_plain.dates.len(), s_imm.dates.len());
    assert_eq!(s_plain.dates[0].day(), 1);
    assert_ne!(s_imm.dates[0].day(), 1);
  }

  #[test]
  fn backward_long_first_merges_short_initial_stub() {
    // Tenor = 2y + 1m → backward semi-annual yields a 1-month short stub at
    // the front. `LongFirst` must absorb that stub into the next regular
    // period, dropping the second date.
    let effective = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
    let termination = NaiveDate::from_ymd_opt(2026, 2, 1).unwrap();
    let short = ScheduleBuilder::new(effective, termination)
      .frequency(Frequency::SemiAnnual)
      .backward()
      .stub(StubConvention::ShortFirst)
      .build();
    let long = ScheduleBuilder::new(effective, termination)
      .frequency(Frequency::SemiAnnual)
      .backward()
      .stub(StubConvention::LongFirst)
      .build();
    assert_eq!(
      long.dates.len() + 1,
      short.dates.len(),
      "LongFirst should drop exactly one intermediate date"
    );
    assert_eq!(long.dates.first(), short.dates.first());
    assert_eq!(long.dates.last(), short.dates.last());
  }

  #[test]
  fn backward_long_first_no_op_on_regular_grid() {
    // Tenor = exact 2y semi-annual → no stub. `LongFirst` must be a no-op.
    let effective = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
    let termination = NaiveDate::from_ymd_opt(2026, 1, 1).unwrap();
    let short = ScheduleBuilder::new(effective, termination)
      .frequency(Frequency::SemiAnnual)
      .backward()
      .stub(StubConvention::ShortFirst)
      .build();
    let long = ScheduleBuilder::new(effective, termination)
      .frequency(Frequency::SemiAnnual)
      .backward()
      .stub(StubConvention::LongFirst)
      .build();
    assert_eq!(
      long.dates, short.dates,
      "LongFirst on a regular tenor must not drop any dates"
    );
  }

  #[test]
  fn forward_long_last_merges_short_trailing_stub() {
    // Tenor = 2y + 1m forward semi-annual → 1-month short stub at the back.
    // `LongLast` must absorb the stub into the previous regular period,
    // dropping `dates[n-2]`.
    let effective = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
    let termination = NaiveDate::from_ymd_opt(2026, 2, 1).unwrap();
    let short = ScheduleBuilder::new(effective, termination)
      .frequency(Frequency::SemiAnnual)
      .forward()
      .stub(StubConvention::ShortLast)
      .build();
    let long = ScheduleBuilder::new(effective, termination)
      .frequency(Frequency::SemiAnnual)
      .forward()
      .stub(StubConvention::LongLast)
      .build();
    assert_eq!(
      long.dates.len() + 1,
      short.dates.len(),
      "LongLast should drop exactly one intermediate date"
    );
    assert_eq!(long.dates.first(), short.dates.first());
    assert_eq!(long.dates.last(), short.dates.last());
  }

  #[test]
  fn forward_long_last_no_op_on_regular_grid() {
    // Exact 1y quarterly grid → no stub. `LongLast` must be a no-op.
    let effective = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
    let termination = NaiveDate::from_ymd_opt(2025, 1, 1).unwrap();
    let short = ScheduleBuilder::new(effective, termination)
      .frequency(Frequency::Quarterly)
      .forward()
      .stub(StubConvention::ShortLast)
      .build();
    let long = ScheduleBuilder::new(effective, termination)
      .frequency(Frequency::Quarterly)
      .forward()
      .stub(StubConvention::LongLast)
      .build();
    assert_eq!(
      long.dates, short.dates,
      "LongLast on a regular tenor must not drop any dates"
    );
  }
}