nyxs_owl 0.5.0

//! # Forecast Trade Module
//!
//! Time series forecasting for financial data using the OxiDiviner library.
//! Provides easy-to-use functions for forecasting stock prices, trading volumes,
//! and other financial time series data.

use chrono::{DateTime, Utc};

/// Re-export key types from oxidiviner for convenience
pub use oxidiviner::{
    api::{ForecastConfig, ForecastOutput, Forecaster, ModelType},
    models::{
        autoregressive::{ARIMAModel, ARModel, SARIMAModel, VARModel},
        exponential_smoothing::{ETSModel, HoltLinearModel, HoltWintersModel, SimpleESModel},
        garch::{EGARCHModel, GARCHModel},
        moving_average::MAModel,
    },
    quick,
};

/// Errors that can occur during forecasting operations
#[derive(Debug, thiserror::Error)]
pub enum ForecastError {
    /// Invalid input data
    #[error("Invalid data: {0}")]
    InvalidData(String),

    /// Invalid parameters
    #[error("Invalid parameters: {0}")]
    InvalidParameters(String),

    /// Model fitting error
    #[error("Model fitting failed: {0}")]
    ModelFittingError(String),

    /// Forecasting error
    #[error("Forecasting failed: {0}")]
    ForecastingError(String),

    /// OxiDiviner error
    #[error("OxiDiviner error: {0}")]
    OxiDivinerError(String),
}

/// Result type for forecasting operations
pub type ForecastResult<T> = std::result::Result<T, ForecastError>;

/// Time series data structure
#[derive(Debug, Clone)]
pub struct TimeSeriesData {
    /// Timestamps for each data point
    pub timestamps: Vec<DateTime<Utc>>,
    /// Values corresponding to each timestamp
    pub values: Vec<f64>,
}

impl TimeSeriesData {
    /// Create new time series data
    pub fn new(timestamps: Vec<DateTime<Utc>>, values: Vec<f64>) -> ForecastResult<Self> {
        if timestamps.len() != values.len() {
            return Err(ForecastError::InvalidData(
                "Timestamps and values must have the same length".to_string(),
            ));
        }

        if timestamps.is_empty() {
            return Err(ForecastError::InvalidData(
                "Time series data cannot be empty".to_string(),
            ));
        }

        // Validate that values are finite
        if values.iter().any(|&v| !v.is_finite()) {
            return Err(ForecastError::InvalidData(
                "All values must be finite numbers".to_string(),
            ));
        }

        Ok(Self { timestamps, values })
    }

    /// Get the length of the time series
    pub fn len(&self) -> usize {
        self.values.len()
    }

    /// Check if the time series is empty
    pub fn is_empty(&self) -> bool {
        self.values.is_empty()
    }

    /// Get the last value in the series
    pub fn last_value(&self) -> Option<f64> {
        self.values.last().copied()
    }

    /// Get a slice of the most recent n values
    pub fn recent_values(&self, n: usize) -> &[f64] {
        let start = self.values.len().saturating_sub(n);
        &self.values[start..]
    }
}

/// Forecast model trait
pub trait ForecastModel {
    /// Fit the model to the given time series data
    fn fit(&mut self, data: &TimeSeriesData) -> ForecastResult<()>;

    /// Generate forecast for the specified number of periods
    fn forecast(&self, periods: usize) -> ForecastResult<Vec<f64>>;

    /// Get the model name
    fn name(&self) -> &str;
}

/// Moving Average Forecasting
///
/// Uses simple moving average to forecast future values
pub fn forecast_moving_average(
    data: &[f64],
    window: usize,
    periods: usize,
) -> ForecastResult<Vec<f64>> {
    if data.is_empty() {
        return Err(ForecastError::InvalidData(
            "Data cannot be empty".to_string(),
        ));
    }

    if window == 0 || periods == 0 {
        return Err(ForecastError::InvalidParameters(
            "Window and periods must be greater than 0".to_string(),
        ));
    }

    if window > data.len() {
        return Err(ForecastError::InvalidParameters(
            "Window size cannot be larger than data length".to_string(),
        ));
    }

    // Create timestamps (using index-based for simplicity)
    let now = Utc::now();
    let timestamps: Vec<DateTime<Utc>> = (0..data.len())
        .map(|i| now + chrono::Duration::days(i as i64))
        .collect();

    // Use oxidiviner's moving average forecast
    quick::ma_forecast(timestamps, data.to_vec(), periods)
        .map_err(|e| ForecastError::OxiDivinerError(e.to_string()))
}

/// Exponential Smoothing Forecasting
///
/// Uses exponential smoothing to forecast future values
pub fn forecast_exponential_smoothing(
    data: &[f64],
    alpha: f64,
    periods: usize,
) -> ForecastResult<Vec<f64>> {
    if data.is_empty() {
        return Err(ForecastError::InvalidData(
            "Data cannot be empty".to_string(),
        ));
    }

    if !(0.0 < alpha && alpha < 1.0) {
        return Err(ForecastError::InvalidParameters(
            "Alpha must be between 0 and 1 (exclusive)".to_string(),
        ));
    }

    if periods == 0 {
        return Err(ForecastError::InvalidParameters(
            "Periods must be greater than 0".to_string(),
        ));
    }

    // Create timestamps
    let now = Utc::now();
    let timestamps: Vec<DateTime<Utc>> = (0..data.len())
        .map(|i| now + chrono::Duration::days(i as i64))
        .collect();

    // Use oxidiviner's exponential smoothing forecast
    quick::es_forecast(timestamps, data.to_vec(), periods)
        .map_err(|e| ForecastError::OxiDivinerError(e.to_string()))
}

/// ARIMA Forecasting
///
/// Uses ARIMA model to forecast future values
pub fn forecast_arima(
    data: &[f64],
    order: (usize, usize, usize),
    periods: usize,
) -> ForecastResult<Vec<f64>> {
    if data.is_empty() {
        return Err(ForecastError::InvalidData(
            "Data cannot be empty".to_string(),
        ));
    }

    if periods == 0 {
        return Err(ForecastError::InvalidParameters(
            "Periods must be greater than 0".to_string(),
        ));
    }

    let (p, d, q) = order;
    if p + d + q > data.len() / 2 {
        return Err(ForecastError::InvalidParameters(
            "Model order too high for data length".to_string(),
        ));
    }

    // Create timestamps
    let now = Utc::now();
    let timestamps: Vec<DateTime<Utc>> = (0..data.len())
        .map(|i| now + chrono::Duration::days(i as i64))
        .collect();

    // Use oxidiviner's ARIMA forecast
    quick::arima_forecast(timestamps, data.to_vec(), periods)
        .map_err(|e| ForecastError::OxiDivinerError(e.to_string()))
}

/// Easy API module for quick forecasting
pub mod easy {
    use super::*;

    /// Automatic model selection and forecasting
    ///
    /// Tries different models and returns the first successful forecast
    pub fn auto_forecast(
        timestamps: &[DateTime<Utc>],
        values: &[f64],
        periods: usize,
    ) -> ForecastResult<(Vec<f64>, String)> {
        if timestamps.len() != values.len() {
            return Err(ForecastError::InvalidData(
                "Timestamps and values must have the same length".to_string(),
            ));
        }

        // Use oxidiviner's auto_forecast which handles model selection
        quick::auto_forecast(timestamps.to_vec(), values.to_vec(), periods)
            .map_err(|e| ForecastError::OxiDivinerError(e.to_string()))
    }

    /// Quick moving average forecast with default window
    pub fn ma_forecast(
        timestamps: &[DateTime<Utc>],
        values: &[f64],
        periods: usize,
    ) -> ForecastResult<Vec<f64>> {
        quick::ma_forecast(timestamps.to_vec(), values.to_vec(), periods)
            .map_err(|e| ForecastError::OxiDivinerError(e.to_string()))
    }

    /// Quick exponential smoothing forecast
    pub fn es_forecast(
        timestamps: &[DateTime<Utc>],
        values: &[f64],
        periods: usize,
    ) -> ForecastResult<Vec<f64>> {
        quick::es_forecast(timestamps.to_vec(), values.to_vec(), periods)
            .map_err(|e| ForecastError::OxiDivinerError(e.to_string()))
    }

    /// Quick ARIMA forecast
    pub fn arima_forecast(
        timestamps: &[DateTime<Utc>],
        values: &[f64],
        periods: usize,
    ) -> ForecastResult<Vec<f64>> {
        quick::arima_forecast(timestamps.to_vec(), values.to_vec(), periods)
            .map_err(|e| ForecastError::OxiDivinerError(e.to_string()))
    }

    /// Compare multiple models and return the best one
    pub fn model_comparison(
        timestamps: &[DateTime<Utc>],
        values: &[f64],
        periods: usize,
    ) -> ForecastResult<Vec<(String, Vec<f64>)>> {
        let mut results = Vec::new();

        // Try different forecasting methods
        if let Ok(forecast) = ma_forecast(timestamps, values, periods) {
            results.push(("Moving Average".to_string(), forecast));
        }

        if let Ok(forecast) = es_forecast(timestamps, values, periods) {
            results.push(("Exponential Smoothing".to_string(), forecast));
        }

        if let Ok(forecast) = arima_forecast(timestamps, values, periods) {
            results.push(("ARIMA".to_string(), forecast));
        }

        if results.is_empty() {
            Err(ForecastError::ModelFittingError(
                "No models could successfully generate a forecast".to_string(),
            ))
        } else {
            Ok(results)
        }
    }

    /// Quick forecast for financial data (daily prices)
    pub fn financial_forecast(
        prices: &[f64],
        periods: usize,
    ) -> ForecastResult<(Vec<f64>, String)> {
        if prices.is_empty() {
            return Err(ForecastError::InvalidData(
                "Price data cannot be empty".to_string(),
            ));
        }

        // Generate daily timestamps (assuming daily data)
        let now = Utc::now();
        let timestamps: Vec<DateTime<Utc>> = (0..prices.len())
            .map(|i| now + chrono::Duration::days(i as i64))
            .collect();

        auto_forecast(&timestamps, prices, periods)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Duration;

    fn create_test_data() -> Vec<f64> {
        (0..30)
            .map(|i| 100.0 + i as f64 + (i as f64 * 0.1).sin() * 5.0)
            .collect()
    }

    fn create_test_timestamps(len: usize) -> Vec<DateTime<Utc>> {
        let start = Utc::now();
        (0..len).map(|i| start + Duration::days(i as i64)).collect()
    }

    #[test]
    fn test_time_series_data_creation() {
        let timestamps = create_test_timestamps(10);
        let values = create_test_data()[..10].to_vec();

        let ts_data = TimeSeriesData::new(timestamps, values);
        assert!(ts_data.is_ok());

        let ts_data = ts_data.unwrap();
        assert_eq!(ts_data.len(), 10);
        assert!(!ts_data.is_empty());
    }

    #[test]
    fn test_time_series_data_validation() {
        let timestamps = create_test_timestamps(5);
        let values = vec![1.0, 2.0, 3.0]; // Wrong length

        let result = TimeSeriesData::new(timestamps, values);
        assert!(result.is_err());
    }

    #[test]
    fn test_moving_average_forecast() {
        let data = create_test_data();
        let result = forecast_moving_average(&data, 5, 3);

        match result {
            Ok(forecast) => {
                assert_eq!(forecast.len(), 3);
                assert!(forecast.iter().all(|&x| x.is_finite()));
            }
            Err(_) => {
                // It's okay if this fails since it depends on external library
                println!("Moving average forecast failed (expected in some test environments)");
            }
        }
    }

    #[test]
    fn test_exponential_smoothing_forecast() {
        let data = create_test_data();
        let result = forecast_exponential_smoothing(&data, 0.3, 3);

        match result {
            Ok(forecast) => {
                assert_eq!(forecast.len(), 3);
                assert!(forecast.iter().all(|&x| x.is_finite()));
            }
            Err(_) => {
                println!(
                    "Exponential smoothing forecast failed (expected in some test environments)"
                );
            }
        }
    }

    #[test]
    fn test_error_conditions() {
        // Test empty data
        assert!(forecast_moving_average(&[], 5, 3).is_err());

        // Test invalid parameters
        let data = create_test_data();
        assert!(forecast_moving_average(&data, 0, 3).is_err());
        assert!(forecast_moving_average(&data, 100, 3).is_err());
        assert!(forecast_exponential_smoothing(&data, 0.0, 3).is_err());
        assert!(forecast_exponential_smoothing(&data, 1.0, 3).is_err());
    }
}