//! Validation tests comparing Rust implementation against Nixtla statsforecast.
//!
//! These tests load reference outputs generated by scripts/generate_nixtla_reference.py
//! and compare them against our Rust implementation.

use anofox_forecast::core::{CalendarAnnotations, TimeSeries};
use anofox_forecast::models::arima::AutoARIMA;
use anofox_forecast::models::mfles::MFLES;
use anofox_forecast::models::Forecaster;
use chrono::{NaiveDate, TimeZone, Utc};
use serde::Deserialize;
use std::collections::HashMap;
use std::fs;

#[derive(Deserialize)]
struct TestData {
    dates: Vec<String>,
    y: Vec<f64>,
    x1: Vec<f64>,
    x2: Vec<f64>,
    true_coefficients: TrueCoefficients,
}

#[derive(Deserialize)]
struct TrueCoefficients {
    intercept: f64,
    x1: f64,
    x2: f64,
}

#[derive(Deserialize)]
struct ReferenceOutput {
    #[allow(dead_code)]
    model: String,
    horizon: usize,
    forecast: Vec<f64>,
    future_x1: Vec<f64>,
    future_x2: Vec<f64>,
}

fn load_test_data() -> TestData {
    let data = fs::read_to_string("tests/reference/test_data_exog.json")
        .expect("Failed to read test data. Run scripts/generate_nixtla_reference.py first.");
    serde_json::from_str(&data).expect("Failed to parse test data")
}

fn load_reference(model: &str) -> ReferenceOutput {
    let filename = format!("tests/reference/{}_exog_reference.json", model);
    let data = fs::read_to_string(&filename).unwrap_or_else(|_| {
        panic!(
            "Failed to read {}. Run scripts/generate_nixtla_reference.py first.",
            filename
        )
    });
    serde_json::from_str(&data).expect("Failed to parse reference data")
}

fn create_time_series(data: &TestData) -> TimeSeries {
    let timestamps: Vec<_> = data
        .dates
        .iter()
        .map(|d| {
            let date = NaiveDate::parse_from_str(d, "%Y-%m-%d").unwrap();
            Utc.from_utc_datetime(&date.and_hms_opt(0, 0, 0).unwrap())
        })
        .collect();

    let calendar = CalendarAnnotations::new()
        .with_regressor("x1".to_string(), data.x1.clone())
        .with_regressor("x2".to_string(), data.x2.clone());

    let mut ts = TimeSeries::univariate(timestamps, data.y.clone()).unwrap();
    ts.set_calendar(calendar);
    ts
}

fn create_future_regressors(reference: &ReferenceOutput) -> HashMap<String, Vec<f64>> {
    let mut future = HashMap::new();
    future.insert("x1".to_string(), reference.future_x1.clone());
    future.insert("x2".to_string(), reference.future_x2.clone());
    future
}

/// Calculate Mean Absolute Percentage Error between two forecast vectors.
fn mape(actual: &[f64], predicted: &[f64]) -> f64 {
    if actual.len() != predicted.len() || actual.is_empty() {
        return f64::MAX;
    }

    let sum: f64 = actual
        .iter()
        .zip(predicted.iter())
        .map(|(a, p)| ((a - p) / a).abs())
        .sum();

    100.0 * sum / actual.len() as f64
}

/// Calculate Root Mean Squared Error between two forecast vectors.
fn rmse(actual: &[f64], predicted: &[f64]) -> f64 {
    if actual.len() != predicted.len() || actual.is_empty() {
        return f64::MAX;
    }

    let sum: f64 = actual
        .iter()
        .zip(predicted.iter())
        .map(|(a, p)| (a - p).powi(2))
        .sum();

    (sum / actual.len() as f64).sqrt()
}

#[test]
fn validate_auto_arima_against_nixtla() {
    let data = load_test_data();
    let reference = load_reference("arima");
    let ts = create_time_series(&data);
    let future = create_future_regressors(&reference);

    let mut model = AutoARIMA::new();
    model.fit(&ts).unwrap();

    let forecast = model.predict_with_exog(reference.horizon, &future).unwrap();
    let rust_forecast = forecast.primary();

    // Calculate error metrics
    let error = mape(&reference.forecast, rust_forecast);
    let rmse_val = rmse(&reference.forecast, rust_forecast);

    println!("AutoARIMA Validation:");
    println!("  Nixtla forecast: {:?}", &reference.forecast[..3]);
    println!("  Rust forecast:   {:?}", &rust_forecast[..3]);
    println!("  MAPE: {:.2}%", error);
    println!("  RMSE: {:.4}", rmse_val);

    // Allow up to 15% MAPE difference (ARIMA model selection may vary)
    // Different ARIMA order selection will lead to different forecasts
    assert!(
        error < 15.0,
        "AutoARIMA MAPE {:.2}% exceeds 15% threshold. Nixtla: {:?}, Rust: {:?}",
        error,
        reference.forecast,
        rust_forecast
    );
}

#[test]
fn validate_mfles_against_nixtla() {
    let data = load_test_data();
    let reference = load_reference("mfles");
    let ts = create_time_series(&data);
    let future = create_future_regressors(&reference);

    let mut model = MFLES::new(vec![7]); // Weekly seasonality
    model.fit(&ts).unwrap();

    let forecast = model.predict_with_exog(reference.horizon, &future).unwrap();
    let rust_forecast = forecast.primary();

    // Calculate error metrics
    let error = mape(&reference.forecast, rust_forecast);
    let rmse_val = rmse(&reference.forecast, rust_forecast);

    println!("MFLES Validation:");
    println!("  Nixtla forecast: {:?}", &reference.forecast[..3]);
    println!("  Rust forecast:   {:?}", &rust_forecast[..3]);
    println!("  MAPE: {:.2}%", error);
    println!("  RMSE: {:.4}", rmse_val);

    // Allow up to 10% MAPE difference for MFLES
    assert!(
        error < 10.0,
        "MFLES MAPE {:.2}% exceeds 10% threshold. Nixtla: {:?}, Rust: {:?}",
        error,
        reference.forecast,
        rust_forecast
    );
}

#[test]
fn validate_exogenous_coefficients_recovery() {
    // Test that OLS correctly recovers the true coefficients
    let data = load_test_data();
    let ts = create_time_series(&data);

    // Use ARIMA(0,0,0) which is pure regression
    use anofox_forecast::models::arima::ARIMA;
    let mut model = ARIMA::new(0, 0, 0);
    model.fit(&ts).unwrap();

    // The true coefficients are:
    // intercept = 50.0, x1 = 5.0, x2 = 10.0
    // With AR component in data, we won't get exact recovery,
    // but should be in the ballpark

    println!("True coefficients:");
    println!("  intercept: {}", data.true_coefficients.intercept);
    println!("  x1: {}", data.true_coefficients.x1);
    println!("  x2: {}", data.true_coefficients.x2);

    // Model should have detected exogenous regressors
    assert!(model.has_exog());
    assert_eq!(model.exog_names().unwrap().len(), 2);
}

#[test]
fn validate_forecast_direction_matches_nixtla() {
    // Even if exact values differ, forecast direction should match
    let data = load_test_data();
    let reference = load_reference("arima");
    let ts = create_time_series(&data);
    let future = create_future_regressors(&reference);

    let mut model = AutoARIMA::new();
    model.fit(&ts).unwrap();

    let forecast = model.predict_with_exog(reference.horizon, &future).unwrap();
    let rust_forecast = forecast.primary();

    // Check that forecast direction (up/down) matches for most points
    let mut direction_matches = 0;
    for i in 1..reference.horizon {
        let nixtla_up = reference.forecast[i] > reference.forecast[i - 1];
        let rust_up = rust_forecast[i] > rust_forecast[i - 1];
        if nixtla_up == rust_up {
            direction_matches += 1;
        }
    }

    let direction_accuracy = direction_matches as f64 / (reference.horizon - 1) as f64;
    println!(
        "Forecast direction accuracy: {:.1}% ({}/{})",
        direction_accuracy * 100.0,
        direction_matches,
        reference.horizon - 1
    );

    // At least 60% of directions should match
    assert!(
        direction_accuracy >= 0.6,
        "Forecast direction accuracy {:.1}% is below 60% threshold",
        direction_accuracy * 100.0
    );
}
