use crate::FinDate;
use chrono::{Datelike, Days, Months, NaiveDate};
use crate::algebra::{self, adjust, checked_add_years};
use crate::calendar::Calendar;
use crate::conventions::{AdjustRule, Frequency};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Schedule<'a> {
pub frequency: Frequency,
pub calendar: Option<&'a Calendar>,
pub adjust_rule: Option<AdjustRule>,
}
impl<'a> Schedule<'a> {
pub fn new(
frequency: Frequency,
opt_calendar: Option<&'a Calendar>,
opt_adjust_rule: Option<AdjustRule>,
) -> Self {
Self {
frequency,
calendar: opt_calendar,
adjust_rule: opt_adjust_rule,
}
}
pub fn iter(&self, anchor: FinDate) -> ScheduleIterator<'_> {
ScheduleIterator {
schedule: self,
anchor,
}
}
pub fn generate(
&self,
anchor_date: &FinDate,
end_date: &FinDate,
) -> Result<Vec<FinDate>, &'static str> {
if end_date <= anchor_date {
return Err("Anchor date must be before end date");
}
if self.frequency == Frequency::Zero {
let adjusted_end = adjust(end_date, self.calendar, self.adjust_rule);
return Ok(vec![adjusted_end]);
}
let mut res = vec![adjust(anchor_date, self.calendar, self.adjust_rule)];
let mut current = *anchor_date;
while let Some(next) = schedule_next(¤t, self.frequency) {
if next > *end_date {
break;
}
res.push(adjust(&next, self.calendar, self.adjust_rule));
current = next;
}
res.dedup();
Ok(res)
}
}
fn force_adjust(
anchor_date: &FinDate,
next_date: &FinDate,
opt_calendar: Option<&Calendar>,
opt_adjust_rule: Option<AdjustRule>,
) -> Option<FinDate> {
let mut res = algebra::adjust(next_date, opt_calendar, opt_adjust_rule);
let mut day_i = 1u64;
while res <= *anchor_date {
let candidate = next_date.checked_add_days(Days::new(day_i))?;
res = algebra::adjust(&candidate, opt_calendar, opt_adjust_rule);
day_i += 1;
}
Some(res)
}
fn schedule_next(anchor_date: &FinDate, frequency: Frequency) -> Option<FinDate> {
match frequency {
Frequency::Daily => anchor_date.checked_add_days(Days::new(1)),
Frequency::Weekly => anchor_date.checked_add_days(Days::new(7)),
Frequency::Biweekly => anchor_date.checked_add_days(Days::new(14)),
Frequency::EveryFourthWeek => anchor_date.checked_add_days(Days::new(28)),
Frequency::Monthly => anchor_date.checked_add_months(Months::new(1)),
Frequency::EndOfMonth => {
let next = anchor_date.checked_add_months(Months::new(1))?;
let first_of_next = if next.month() == 12 {
NaiveDate::from_ymd_opt(next.year() + 1, 1, 1)
} else {
NaiveDate::from_ymd_opt(next.year(), next.month() + 1, 1)
};
first_of_next.and_then(|d| d.pred_opt())
}
Frequency::Bimonthly => anchor_date.checked_add_months(Months::new(2)),
Frequency::Quarterly => anchor_date.checked_add_months(Months::new(3)),
Frequency::EveryFourthMonth => anchor_date.checked_add_months(Months::new(4)),
Frequency::Semiannual => anchor_date.checked_add_months(Months::new(6)),
Frequency::Annual => checked_add_years(anchor_date, 1),
Frequency::Zero => None,
}
}
pub fn schedule_next_adjusted(schedule: &Schedule, anchor: FinDate) -> Option<FinDate> {
let next = schedule_next(&anchor, schedule.frequency)?;
force_adjust(&anchor, &next, schedule.calendar, schedule.adjust_rule)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ScheduleIterator<'a> {
schedule: &'a Schedule<'a>,
anchor: FinDate,
}
impl<'a> Iterator for ScheduleIterator<'a> {
type Item = FinDate;
fn next(&mut self) -> Option<Self::Item> {
let res = schedule_next_adjusted(self.schedule, self.anchor)?;
self.anchor = res;
Some(res)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn end_of_month_schedule_next_test() {
let sched = Schedule::new(Frequency::EndOfMonth, None, None);
let d = NaiveDate::from_ymd_opt(2024, 1, 31).unwrap();
assert_eq!(schedule_next_adjusted(&sched, d), NaiveDate::from_ymd_opt(2024, 2, 29));
let d = NaiveDate::from_ymd_opt(2023, 1, 31).unwrap();
assert_eq!(schedule_next_adjusted(&sched, d), NaiveDate::from_ymd_opt(2023, 2, 28));
let d = NaiveDate::from_ymd_opt(2024, 11, 30).unwrap();
assert_eq!(schedule_next_adjusted(&sched, d), NaiveDate::from_ymd_opt(2024, 12, 31));
let d = NaiveDate::from_ymd_opt(2024, 1, 15).unwrap();
assert_eq!(schedule_next_adjusted(&sched, d), NaiveDate::from_ymd_opt(2024, 2, 29));
}
}