use chrono::{DateTime, Datelike, Duration, Timelike, Utc};
use shape_ast::ast::{Timeframe, TimeframeUnit};
use shape_ast::error::{Result, ShapeError};
pub fn parse_timeframe_string(s: &str) -> Result<Timeframe> {
let s = s.trim();
if s.is_empty() {
return Err(ShapeError::RuntimeError {
message: "Empty timeframe string".to_string(),
location: None,
});
}
let split_idx = s.find(|c: char| !c.is_ascii_digit()).unwrap_or(s.len());
let (num_str, unit_str) = s.split_at(split_idx);
let value: u32 = num_str.parse().map_err(|_| ShapeError::RuntimeError {
message: format!("Invalid timeframe number: {}", num_str),
location: None,
})?;
let unit = match unit_str.to_lowercase().as_str() {
"m" | "min" => TimeframeUnit::Minute,
"h" | "hour" => TimeframeUnit::Hour,
"d" | "day" => TimeframeUnit::Day,
"w" | "week" => TimeframeUnit::Week,
"mo" | "month" => TimeframeUnit::Month,
"y" | "year" => TimeframeUnit::Year,
"" => TimeframeUnit::Minute, _ => {
return Err(ShapeError::RuntimeError {
message: format!("Unknown timeframe unit: {}", unit_str),
location: None,
});
}
};
Ok(Timeframe::new(value, unit))
}
pub fn timeframe_to_duration(tf: &Timeframe) -> Duration {
Duration::seconds(tf.to_seconds() as i64)
}
pub fn timeframe_to_minutes(tf: &Timeframe) -> i64 {
tf.to_seconds() as i64 / 60
}
pub fn align_timestamp(ts: DateTime<Utc>, tf: &Timeframe) -> DateTime<Utc> {
let seconds = tf.to_seconds() as i64;
if seconds == 60 {
ts.with_second(0).unwrap().with_nanosecond(0).unwrap()
} else if seconds == 300 {
let minute = ts.minute();
let aligned_minute = (minute / 5) * 5;
ts.with_minute(aligned_minute)
.unwrap()
.with_second(0)
.unwrap()
.with_nanosecond(0)
.unwrap()
} else if seconds == 900 {
let minute = ts.minute();
let aligned_minute = (minute / 15) * 15;
ts.with_minute(aligned_minute)
.unwrap()
.with_second(0)
.unwrap()
.with_nanosecond(0)
.unwrap()
} else if seconds == 1800 {
let minute = ts.minute();
let aligned_minute = (minute / 30) * 30;
ts.with_minute(aligned_minute)
.unwrap()
.with_second(0)
.unwrap()
.with_nanosecond(0)
.unwrap()
} else if seconds == 3600 {
ts.with_minute(0)
.unwrap()
.with_second(0)
.unwrap()
.with_nanosecond(0)
.unwrap()
} else if seconds == 14400 {
let hour = ts.hour();
let aligned_hour = (hour / 4) * 4;
ts.with_hour(aligned_hour)
.unwrap()
.with_minute(0)
.unwrap()
.with_second(0)
.unwrap()
.with_nanosecond(0)
.unwrap()
} else if seconds == 86400 {
ts.with_hour(0)
.unwrap()
.with_minute(0)
.unwrap()
.with_second(0)
.unwrap()
.with_nanosecond(0)
.unwrap()
} else if seconds == 604800 {
let days_from_monday = ts.weekday().num_days_from_monday();
let aligned = ts - Duration::days(days_from_monday as i64);
aligned
.with_hour(0)
.unwrap()
.with_minute(0)
.unwrap()
.with_second(0)
.unwrap()
.with_nanosecond(0)
.unwrap()
} else {
let ts_seconds = ts.timestamp();
let aligned_seconds = (ts_seconds / seconds) * seconds;
DateTime::from_timestamp(aligned_seconds, 0).unwrap()
}
}
pub fn generate_timestamps(
start: DateTime<Utc>,
end: DateTime<Utc>,
tf: &Timeframe,
) -> Vec<DateTime<Utc>> {
let mut timestamps = Vec::new();
let duration = timeframe_to_duration(tf);
let mut current = align_timestamp(start, tf);
while current <= end {
timestamps.push(current);
current += duration;
}
timestamps
}
pub fn generate_timestamps_micros(
start: DateTime<Utc>,
end: DateTime<Utc>,
tf: &Timeframe,
) -> Vec<i64> {
generate_timestamps(start, end, tf)
.into_iter()
.map(|ts| ts.timestamp_micros())
.collect()
}
pub fn count_rows(start: DateTime<Utc>, end: DateTime<Utc>, tf: &Timeframe) -> usize {
let duration = end - start;
let tf_duration = timeframe_to_duration(tf);
((duration.num_milliseconds() / tf_duration.num_milliseconds()) + 1) as usize
}
pub fn find_common_timeframe(tf1: &Timeframe, tf2: &Timeframe) -> Timeframe {
let minutes1 = timeframe_to_minutes(tf1);
let minutes2 = timeframe_to_minutes(tf2);
if minutes1 <= minutes2 { *tf1 } else { *tf2 }
}
pub fn is_timeframe_compatible(base: &Timeframe, target: &Timeframe) -> bool {
let base_minutes = timeframe_to_minutes(base);
let target_minutes = timeframe_to_minutes(target);
target_minutes % base_minutes == 0 || base_minutes % target_minutes == 0
}
pub fn find_covering_index(source_timestamps: &[i64], target_timestamp: i64) -> Option<usize> {
if source_timestamps.is_empty() {
return None;
}
match source_timestamps.binary_search(&target_timestamp) {
Ok(idx) => Some(idx),
Err(idx) => {
if idx == 0 {
None
} else {
Some(idx - 1)
}
}
}
}
pub fn alignment_ratio(from_tf: &Timeframe, to_tf: &Timeframe) -> f64 {
let from_minutes = timeframe_to_minutes(from_tf) as f64;
let to_minutes = timeframe_to_minutes(to_tf) as f64;
from_minutes / to_minutes
}
pub fn find_closest_index(timestamps: &[i64], target: i64) -> Option<usize> {
if timestamps.is_empty() {
return None;
}
match timestamps.binary_search(&target) {
Ok(idx) => Some(idx),
Err(idx) => {
if idx == 0 {
Some(0)
} else if idx >= timestamps.len() {
Some(timestamps.len() - 1)
} else {
let diff_prev = (target - timestamps[idx - 1]).abs();
let diff_next = (timestamps[idx] - target).abs();
if diff_prev <= diff_next {
Some(idx - 1)
} else {
Some(idx)
}
}
}
}
}