raphtory-core 0.18.5

use chrono::{DateTime, Datelike, Duration, Months, NaiveDate};
use itertools::Itertools;
use raphtory_api::core::{storage::timeindex::EventTime, utils::time::ParseTimeError};
use regex::Regex;
use std::ops::{Add, Mul, Sub};

pub(crate) const SECOND_MS: i64 = 1000;
pub(crate) const MINUTE_MS: i64 = 60 * SECOND_MS;
pub(crate) const HOUR_MS: i64 = 60 * MINUTE_MS;
pub(crate) const DAY_MS: i64 = 24 * HOUR_MS;
pub(crate) const WEEK_MS: i64 = 7 * DAY_MS;

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum IntervalSize {
    Discrete(u64),
    /// `months` is u32 because chrono::Months works with u32.
    Temporal {
        millis: u64,
        months: u32,
    },
}

impl IntervalSize {
    /// Creates a 0 sized temporal `IntervalSize`. Should not be used to create Windows.
    pub fn empty_temporal() -> Self {
        IntervalSize::Temporal {
            millis: 0,
            months: 0,
        }
    }

    fn months(months: i64) -> Self {
        Self::Temporal {
            millis: 0,
            months: months as u32,
        }
    }

    fn add_temporal(&self, other: IntervalSize) -> IntervalSize {
        match (self, other) {
            (
                Self::Temporal {
                    millis: ml1,
                    months: mt1,
                },
                Self::Temporal {
                    millis: ml2,
                    months: mt2,
                },
            ) => Self::Temporal {
                millis: ml1 + ml2,
                months: mt1 + mt2,
            },
            _ => panic!("this function is not supposed to be used with discrete intervals"),
        }
    }
}

impl From<Duration> for IntervalSize {
    fn from(value: Duration) -> Self {
        Self::Temporal {
            millis: value.num_milliseconds() as u64,
            months: 0,
        }
    }
}

/// Used to keep track of the smallest unit provided, so that we can line up windows
/// (eg. at the start of the hour, week, year, ...)
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum AlignmentUnit {
    Unaligned, // note that there is no functional difference between millisecond and unaligned for the time being
    Millisecond,
    Second,
    Minute,
    Hour,
    Day,
    Week,
    Month,
    Year,
}

impl AlignmentUnit {
    /// Floors a UTC timestamp in milliseconds since Unix epoch to the nearest alignment unit.
    pub fn align_timestamp(&self, timestamp: i64) -> i64 {
        match self {
            AlignmentUnit::Unaligned => timestamp,
            AlignmentUnit::Millisecond => timestamp,
            AlignmentUnit::Second => Self::floor_ms(timestamp, SECOND_MS),
            AlignmentUnit::Minute => Self::floor_ms(timestamp, MINUTE_MS),
            AlignmentUnit::Hour => Self::floor_ms(timestamp, HOUR_MS),
            AlignmentUnit::Day => Self::floor_ms(timestamp, DAY_MS),
            AlignmentUnit::Week => {
                let offset = DAY_MS * 4; // 0 is a Thursday
                Self::floor_ms(timestamp - offset, WEEK_MS) + offset
            }
            // Month and Year are variable (28, 30, or 31 days / 365 or 366 days so we can't simply use division)
            AlignmentUnit::Month => {
                let naive = DateTime::from_timestamp_millis(timestamp)
                    .unwrap_or_else(|| {
                        panic!("{timestamp} cannot be interpreted as a milliseconds timestamp.")
                    })
                    .naive_utc();
                let y = naive.year();
                let m = naive.month();
                NaiveDate::from_ymd_opt(y, m, 1)
                    .unwrap()
                    .and_hms_milli_opt(0, 0, 0, 0)
                    .unwrap()
                    .and_utc()
                    .timestamp_millis()
            }
            AlignmentUnit::Year => {
                let naive = DateTime::from_timestamp_millis(timestamp)
                    .unwrap_or_else(|| {
                        panic!("{timestamp} cannot be interpreted as a milliseconds timestamp.")
                    })
                    .naive_utc();
                let y = naive.year();
                NaiveDate::from_ymd_opt(y, 1, 1)
                    .unwrap()
                    .and_hms_milli_opt(0, 0, 0, 0)
                    .unwrap()
                    .and_utc()
                    .timestamp_millis()
            }
        }
    }

    /// Floors `ts` to a multiple of `unit_ms` using a remainder that is always non-negative,
    /// so the result is the boundary at or before `ts`, even for negative timestamps.
    #[inline]
    fn floor_ms(ts: i64, unit_ms: i64) -> i64 {
        ts - ts.rem_euclid(unit_ms)
    }
}

impl TryFrom<String> for AlignmentUnit {
    type Error = ParseTimeError;

    fn try_from(value: String) -> Result<Self, Self::Error> {
        Self::try_from(value.as_str())
    }
}

impl TryFrom<&str> for AlignmentUnit {
    type Error = ParseTimeError;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        let unit = match value.to_lowercase().as_str() {
            "year" | "years" => AlignmentUnit::Year,
            "month" | "months" => AlignmentUnit::Month,
            "week" | "weeks" => AlignmentUnit::Week,
            "day" | "days" => AlignmentUnit::Day,
            "hour" | "hours" => AlignmentUnit::Hour,
            "minute" | "minutes" => AlignmentUnit::Minute,
            "second" | "seconds" => AlignmentUnit::Second,
            "millisecond" | "milliseconds" => AlignmentUnit::Millisecond,
            "unaligned" => AlignmentUnit::Unaligned,
            unit => return Err(ParseTimeError::InvalidAlignmentUnit(unit.to_string())),
        };
        Ok(unit)
    }
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Interval {
    /// Used if the `IntervalSize` is Temporal, keeps track of the smallest unit passed to line up windows.
    /// (eg. at the start of the hour, week, year, ...). If the IntervalSize is discrete, this is `None`.
    pub alignment_unit: Option<AlignmentUnit>,
    /// The interval.
    pub size: IntervalSize,
}

impl Default for Interval {
    fn default() -> Self {
        Self {
            alignment_unit: None,
            size: IntervalSize::Discrete(1),
        }
    }
}

impl TryFrom<String> for Interval {
    type Error = ParseTimeError;

    fn try_from(value: String) -> Result<Self, Self::Error> {
        Self::try_from(value.as_str())
    }
}

impl TryFrom<&str> for Interval {
    type Error = ParseTimeError;
    fn try_from(value: &str) -> Result<Self, Self::Error> {
        let trimmed = value.trim();
        let no_and = trimmed.replace("and", "");
        let cleaned = {
            let re = Regex::new(r"[\s&,]+").unwrap();
            re.replace_all(&no_and, " ")
        };

        let tokens = cleaned.split(' ').collect_vec();

        if tokens.len() < 2 || tokens.len() % 2 != 0 {
            return Err(ParseTimeError::InvalidPairs);
        }

        let (temporal_sum, smallest_unit): (IntervalSize, AlignmentUnit) =
            tokens.chunks(2).try_fold(
                (IntervalSize::empty_temporal(), AlignmentUnit::Year), // start with the largest alignment unit
                |(sum, smallest), chunk| {
                    let (interval, unit) = Self::parse_duration(chunk[0], chunk[1])?;
                    Ok::<_, ParseTimeError>((sum.add_temporal(interval), smallest.min(unit)))
                },
            )?;

        Ok(Self {
            alignment_unit: Some(smallest_unit),
            size: temporal_sum,
        })
    }
}

impl TryFrom<u64> for Interval {
    type Error = ParseTimeError;
    fn try_from(value: u64) -> Result<Self, Self::Error> {
        Ok(Self {
            alignment_unit: None,
            size: IntervalSize::Discrete(value),
        })
    }
}

impl TryFrom<u32> for Interval {
    type Error = ParseTimeError;
    fn try_from(value: u32) -> Result<Self, Self::Error> {
        Ok(Self {
            alignment_unit: None,
            size: IntervalSize::Discrete(value as u64),
        })
    }
}

impl TryFrom<i32> for Interval {
    type Error = ParseTimeError;
    fn try_from(value: i32) -> Result<Self, Self::Error> {
        if value >= 0 {
            Ok(Self {
                alignment_unit: None,
                size: IntervalSize::Discrete(value as u64),
            })
        } else {
            Err(ParseTimeError::NegativeInt)
        }
    }
}

impl TryFrom<i64> for Interval {
    type Error = ParseTimeError;

    fn try_from(value: i64) -> Result<Self, Self::Error> {
        if value >= 0 {
            Ok(Self {
                alignment_unit: None,
                size: IntervalSize::Discrete(value as u64),
            })
        } else {
            Err(ParseTimeError::NegativeInt)
        }
    }
}

pub trait TryIntoInterval {
    fn try_into_interval(self) -> Result<Interval, ParseTimeError>;
}

impl<T> TryIntoInterval for T
where
    Interval: TryFrom<T>,
    ParseTimeError: From<<Interval as TryFrom<T>>::Error>,
{
    fn try_into_interval(self) -> Result<Interval, ParseTimeError> {
        Ok(self.try_into()?)
    }
}

impl Interval {
    /// Return an option because there might be no exact translation to millis for some intervals
    pub fn to_millis(&self) -> Option<u64> {
        match self.size {
            IntervalSize::Discrete(millis) => Some(millis),
            IntervalSize::Temporal { millis, months } => (months == 0).then_some(millis),
        }
    }

    fn parse_duration(
        number: &str,
        unit: &str,
    ) -> Result<(IntervalSize, AlignmentUnit), ParseTimeError> {
        let number: i64 = number.parse::<u64>()? as i64;
        let duration = match unit {
            "year" | "years" => (IntervalSize::months(number * 12), AlignmentUnit::Year),
            "month" | "months" => (IntervalSize::months(number), AlignmentUnit::Month),
            "week" | "weeks" => (Duration::weeks(number).into(), AlignmentUnit::Week),
            "day" | "days" => (Duration::days(number).into(), AlignmentUnit::Day),
            "hour" | "hours" => (Duration::hours(number).into(), AlignmentUnit::Hour),
            "minute" | "minutes" => (Duration::minutes(number).into(), AlignmentUnit::Minute),
            "second" | "seconds" => (Duration::seconds(number).into(), AlignmentUnit::Second),
            "millisecond" | "milliseconds" => (
                Duration::milliseconds(number).into(),
                AlignmentUnit::Millisecond,
            ),
            unit => return Err(ParseTimeError::InvalidUnit(unit.to_string())),
        };
        Ok(duration)
    }

    pub fn discrete(num: u64) -> Self {
        Interval {
            alignment_unit: None,
            size: IntervalSize::Discrete(num),
        }
    }

    pub fn milliseconds(ms: i64) -> Self {
        Interval {
            alignment_unit: Some(AlignmentUnit::Millisecond),
            size: IntervalSize::from(Duration::milliseconds(ms)),
        }
    }

    pub fn seconds(seconds: i64) -> Self {
        Interval {
            alignment_unit: Some(AlignmentUnit::Second),
            size: IntervalSize::from(Duration::seconds(seconds)),
        }
    }

    pub fn minutes(minutes: i64) -> Self {
        Interval {
            alignment_unit: Some(AlignmentUnit::Minute),
            size: IntervalSize::from(Duration::minutes(minutes)),
        }
    }

    pub fn hours(hours: i64) -> Self {
        Interval {
            alignment_unit: Some(AlignmentUnit::Hour),
            size: IntervalSize::from(Duration::hours(hours)),
        }
    }

    pub fn days(days: i64) -> Self {
        Interval {
            alignment_unit: Some(AlignmentUnit::Day),
            size: IntervalSize::from(Duration::days(days)),
        }
    }

    pub fn weeks(weeks: i64) -> Self {
        Interval {
            alignment_unit: Some(AlignmentUnit::Week),
            size: IntervalSize::from(Duration::weeks(weeks)),
        }
    }

    pub fn months(months: i64) -> Self {
        Interval {
            alignment_unit: Some(AlignmentUnit::Month),
            size: IntervalSize::months(months),
        }
    }

    pub fn years(years: i64) -> Self {
        Interval {
            alignment_unit: Some(AlignmentUnit::Year),
            size: IntervalSize::months(12 * years),
        }
    }

    pub fn and(&self, other: &Self) -> Result<Self, IntervalTypeError> {
        match (self.size, other.size) {
            (IntervalSize::Discrete(l), IntervalSize::Discrete(r)) => Ok(Interval {
                alignment_unit: None,
                size: IntervalSize::Discrete(l + r),
            }),
            (IntervalSize::Temporal { .. }, IntervalSize::Temporal { .. }) => Ok(Interval {
                alignment_unit: self.alignment_unit.min(other.alignment_unit),
                size: self.size.add_temporal(other.size),
            }),
            (_, _) => Err(IntervalTypeError()),
        }
    }
}

#[derive(thiserror::Error, Debug)]
#[error("Discrete and temporal intervals cannot be combined")]
pub struct IntervalTypeError();

impl Sub<Interval> for i64 {
    type Output = i64;
    fn sub(self, rhs: Interval) -> Self::Output {
        match rhs.size {
            IntervalSize::Discrete(number)
            | IntervalSize::Temporal {
                millis: number,
                months: 0,
            } => self - (number as i64),
            IntervalSize::Temporal { millis, months } => {
                // first we subtract the number of milliseconds and then the number of months for
                // consistency with the implementation of Add (we revert back the steps) so we
                // guarantee that:  time + interval - interval = time
                let datetime = DateTime::from_timestamp_millis(self - millis as i64)
                    .unwrap_or_else(|| {
                        panic!("{self} cannot be interpreted as a milliseconds timestamp")
                    })
                    .naive_utc();
                (datetime - Months::new(months))
                    .and_utc()
                    .timestamp_millis()
            }
        }
    }
}

impl Add<Interval> for i64 {
    type Output = i64;
    fn add(self, rhs: Interval) -> Self::Output {
        match rhs.size {
            IntervalSize::Discrete(number)
            | IntervalSize::Temporal {
                millis: number,
                months: 0,
            } => self + (number as i64),
            IntervalSize::Temporal { millis, months } => {
                // first we add the number of months and then the number of milliseconds for
                // consistency with the implementation of Sub (we revert back the steps) so we
                // guarantee that:  time + interval - interval = time
                let datetime = DateTime::from_timestamp_millis(self)
                    .unwrap_or_else(|| {
                        panic!("{self} cannot be interpreted as a milliseconds timestamp")
                    })
                    .naive_utc();
                (datetime + Months::new(months))
                    .and_utc()
                    .timestamp_millis()
                    + millis as i64
            }
        }
    }
}

// since all IntervalSize values (discrete number and temporal millis/months) are unsigned,
// we can only multiply with unsigned numbers.
impl Mul<Interval> for u32 {
    type Output = Interval;

    fn mul(self, rhs: Interval) -> Self::Output {
        match rhs.size {
            IntervalSize::Discrete(number) => Interval {
                alignment_unit: rhs.alignment_unit, // alignment_unit should be None
                size: IntervalSize::Discrete((self as u64) * number),
            },
            IntervalSize::Temporal { millis, months } => Interval {
                alignment_unit: rhs.alignment_unit,
                size: IntervalSize::Temporal {
                    millis: (self as u64) * millis,
                    months: self * months,
                },
            },
        }
    }
}

impl Add<Interval> for EventTime {
    type Output = EventTime;
    fn add(self, rhs: Interval) -> Self::Output {
        match rhs.size {
            IntervalSize::Discrete(number) => EventTime(self.0 + (number as i64), self.1),
            IntervalSize::Temporal { millis, months } => {
                // first we add the number of months and then the number of milliseconds for
                // consistency with the implementation of Sub (we revert back the steps) so we
                // guarantee that:  time + interval - interval = time
                let datetime = DateTime::from_timestamp_millis(self.0)
                    .unwrap_or_else(|| {
                        panic!("{self} cannot be interpreted as a milliseconds timestamp")
                    })
                    .naive_utc();
                let timestamp = (datetime + Months::new(months))
                    .and_utc()
                    .timestamp_millis()
                    + millis as i64;
                EventTime(timestamp, self.1)
            }
        }
    }
}

#[cfg(test)]
mod time_tests {
    use crate::utils::time::{AlignmentUnit, Interval, WEEK_MS};
    use chrono::{DateTime, Datelike, NaiveTime, Utc, Weekday};
    use proptest::{arbitrary::any, prelude::Strategy, proptest};
    use raphtory_api::core::{
        storage::timeindex::AsTime,
        utils::time::{ParseTimeError, TryIntoTime},
    };

    #[test]
    fn alignment_week_proptest() {
        proptest!(|(dt in (-8334601228800000i64..8210266876800000).prop_filter_map("not a valid date", DateTime::from_timestamp_millis))| {
            let ts = dt.timestamp_millis();
            let aligned = AlignmentUnit::Week.align_timestamp(ts);
            let aligned_dt = aligned.dt().unwrap();
            assert_eq!(aligned_dt, aligned_dt.with_time(NaiveTime::from_num_seconds_from_midnight_opt(0, 0).unwrap()).unwrap());
            assert!(ts - aligned < WEEK_MS);
            assert_eq!(aligned_dt.weekday(), Weekday::Mon);


        })
    }

    #[test]
    fn interval_parsing() {
        let second: u64 = 1000;
        let minute = 60 * second;
        let hour = 60 * minute;
        let day = 24 * hour;
        let week = 7 * day;

        let interval: Interval = "1 day".try_into().unwrap();
        assert_eq!(interval.to_millis().unwrap(), day);

        let interval: Interval = "1 week".try_into().unwrap();
        assert_eq!(interval.to_millis().unwrap(), week);

        let interval: Interval = "4 weeks and 1 day".try_into().unwrap();
        assert_eq!(interval.to_millis().unwrap(), 4 * week + day);

        let interval: Interval = "2 days & 1 millisecond".try_into().unwrap();
        assert_eq!(interval.to_millis().unwrap(), 2 * day + 1);

        let interval: Interval = "2 days, 1 hour, and 2 minutes".try_into().unwrap();
        assert_eq!(interval.to_millis().unwrap(), 2 * day + hour + 2 * minute);

        let interval: Interval = "1 weeks ,   1 minute".try_into().unwrap();
        assert_eq!(interval.to_millis().unwrap(), week + minute);

        let interval: Interval = "23 seconds  and 34 millisecond and 1 minute"
            .try_into()
            .unwrap();
        assert_eq!(interval.to_millis().unwrap(), 23 * second + 34 + minute);
    }

    #[test]
    fn interval_parsing_with_months_and_years() {
        let dt = "2020-01-01 00:00:00".try_into_time().unwrap();

        let two_months: Interval = "2 months".try_into().unwrap();
        let dt_plus_2_months = "2020-03-01 00:00:00".try_into_time().unwrap();
        assert_eq!(dt + two_months, dt_plus_2_months);

        let two_years: Interval = "2 years".try_into().unwrap();
        let dt_plus_2_years = "2022-01-01 00:00:00".try_into_time().unwrap();
        assert_eq!(dt + two_years, dt_plus_2_years);

        let mix_interval: Interval = "1 year 1 month and 1 second".try_into().unwrap();
        let dt_mix = "2021-02-01 00:00:01".try_into_time().unwrap();
        assert_eq!(dt + mix_interval, dt_mix);
    }

    #[test]
    fn invalid_intervals() {
        let result: Result<Interval, ParseTimeError> = "".try_into();
        assert_eq!(result, Err(ParseTimeError::InvalidPairs));

        let result: Result<Interval, ParseTimeError> = "1".try_into();
        assert_eq!(result, Err(ParseTimeError::InvalidPairs));

        let result: Result<Interval, ParseTimeError> = "1 day and 5".try_into();
        assert_eq!(result, Err(ParseTimeError::InvalidPairs));

        let result: Result<Interval, ParseTimeError> = "1 daay".try_into();
        assert_eq!(result, Err(ParseTimeError::InvalidUnit("daay".to_string())));

        let result: Result<Interval, ParseTimeError> = "day 1".try_into();

        match result {
            Err(ParseTimeError::ParseInt { .. }) => (),
            _ => panic!(),
        }
    }
}