//! Walk-forward analysis.
//!
//! Splits the data into N consecutive train/test folds, optimizes parameters
//! in-sample (train), then evaluates the best parameters out-of-sample (test).
//! Supports anchored (expanding window) and rolling (fixed window) methods.

use std::collections::HashMap;

use bt_analytics::PerformanceMetrics;
use bt_core::alignment::slice_aligned_bars;
use bt_core::orchestrator::{BacktestConfig, StrategyPlan};
use bt_core::sweep::{expand_param_grid, strategy_with_params, ParamGrid};
use bt_core::{BatchResultLite, SimpleBacktestRunner};
use bt_data::DataStore;
use bt_expr::ScalarValue;
use bt_kernel::{BtError, Result, TimeRange, Timestamp};
use serde::{Deserialize, Serialize};

use crate::extract_metric;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum WalkForwardMethod {
    /// Train window grows from the beginning; test window is fixed size.
    Anchored,
    /// Fixed-size train window slides forward.
    Rolling,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct WalkForwardConfig {
    pub method: WalkForwardMethod,
    pub n_splits: usize,
    pub train_ratio: f64,
    pub optimize_metric: String,
    pub param_grid: ParamGrid,
    pub max_parallelism: usize,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct FoldResult {
    pub fold_index: usize,
    pub train_range: TimeRange,
    pub test_range: TimeRange,
    /// Best in-sample metrics (metrics of best param combo on train set).
    pub is_metrics: PerformanceMetrics,
    /// Out-of-sample metrics (best IS params applied to test set).
    pub oos_metrics: PerformanceMetrics,
    /// Best parameters found during in-sample optimization.
    pub best_params: HashMap<String, ScalarValue>,
    /// All IS results for all parameter combos.
    pub all_is_results: Vec<BatchResultLite>,
    /// IS equity curve (best params on train data). Empty if unavailable.
    #[serde(default)]
    pub is_equity: Vec<f64>,
    /// OOS equity curve (best params on test data). Empty if unavailable.
    #[serde(default)]
    pub oos_equity: Vec<f64>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct WalkForwardResult {
    pub folds: Vec<FoldResult>,
    pub best_params_per_fold: Vec<HashMap<String, ScalarValue>>,
}

/// Run walk-forward analysis.
///
/// 1. Loads data once for the full time range.
/// 2. Splits into N train/test folds.
/// 3. For each fold: optimize params on train, evaluate best on test.
pub fn run_walk_forward(
    strategy: &StrategyPlan,
    wf_config: &WalkForwardConfig,
    backtest_config: &BacktestConfig,
    data_store: &dyn DataStore,
) -> Result<WalkForwardResult> {
    if wf_config.n_splits == 0 {
        return Err(BtError::Strategy(
            "walk_forward: n_splits must be > 0".into(),
        ));
    }
    if wf_config.train_ratio <= 0.0 || wf_config.train_ratio >= 1.0 {
        return Err(BtError::Strategy(
            "walk_forward: train_ratio must be in (0, 1)".into(),
        ));
    }

    // Phase 1: Load and align data once for the full time range.
    let runner = SimpleBacktestRunner::new();
    let (aligned, universe_sorted, _version, _, _) =
        SimpleBacktestRunner::load_and_align(backtest_config, data_store)?;

    // Compute fold boundaries.
    let full_start = backtest_config.time_range.start.as_nanos();
    let full_end = backtest_config.time_range.end.as_nanos();
    let full_duration = full_end - full_start;

    let folds = compute_fold_ranges(
        full_start,
        full_end,
        full_duration,
        wf_config.n_splits,
        wf_config.train_ratio,
        &wf_config.method,
    )?;

    // Expand param grid once.
    let param_combos = expand_param_grid(&wf_config.param_grid);

    let mut fold_results = Vec::with_capacity(folds.len());
    let mut best_params_per_fold = Vec::with_capacity(folds.len());

    for (fold_idx, (train_range, test_range)) in folds.iter().enumerate() {
        tracing::info!(
            fold = fold_idx,
            train_start = train_range.start.as_nanos(),
            train_end = train_range.end.as_nanos(),
            test_start = test_range.start.as_nanos(),
            test_end = test_range.end.as_nanos(),
            "walk_forward: processing fold"
        );

        // Slice aligned bars for train and test.
        let train_aligned = slice_aligned_bars(&aligned, *train_range)?;
        let test_aligned = slice_aligned_bars(&aligned, *test_range)?;

        // Build strategies for all param combos.
        let strategies: Vec<StrategyPlan> = param_combos
            .iter()
            .map(|params| strategy_with_params(strategy, params))
            .collect::<Result<Vec<_>>>()?;

        // --- In-sample: run all combos on train data ---
        let mut train_config = backtest_config.clone();
        train_config.time_range = *train_range;

        let is_results = run_on_pre_aligned(
            &runner,
            &strategies,
            &train_config,
            &train_aligned,
            &universe_sorted,
            wf_config.max_parallelism,
        )?;

        // Find best combo by optimize_metric.
        let (best_idx, _best_metric) = is_results
            .iter()
            .enumerate()
            .map(|(i, r)| (i, extract_metric(&r.metrics, &wf_config.optimize_metric)))
            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
            .ok_or_else(|| BtError::Strategy("walk_forward: no IS results".into()))?;

        let best_params = if best_idx < param_combos.len() {
            param_combos[best_idx].clone()
        } else {
            HashMap::new()
        };

        // --- Full runs for best params (IS + OOS) to get equity curves ---
        let best_strategy = strategy_with_params(strategy, &best_params)?;

        // IS full run (best params on train data)
        let mut train_config = backtest_config.clone();
        train_config.time_range = *train_range;
        let is_full = runner.run_on_aligned(
            &best_strategy,
            &train_config,
            &train_aligned,
            &universe_sorted,
            "walk_forward",
        )?;
        let is_equity = is_full
            .equity_curve
            .as_any()
            .downcast_ref::<arrow::array::Float64Array>()
            .map(|a: &arrow::array::Float64Array| a.values().to_vec())
            .unwrap_or_default();
        let is_best_metrics = is_results[best_idx].metrics.clone();

        // OOS full run (best params on test data)
        let mut test_config = backtest_config.clone();
        test_config.time_range = *test_range;
        let oos_full = runner.run_on_aligned(
            &best_strategy,
            &test_config,
            &test_aligned,
            &universe_sorted,
            "walk_forward",
        )?;
        let oos_equity = oos_full
            .equity_curve
            .as_any()
            .downcast_ref::<arrow::array::Float64Array>()
            .map(|a: &arrow::array::Float64Array| a.values().to_vec())
            .unwrap_or_default();
        let oos_metrics = oos_full.metrics;

        fold_results.push(FoldResult {
            fold_index: fold_idx,
            train_range: *train_range,
            test_range: *test_range,
            is_metrics: is_best_metrics,
            oos_metrics,
            best_params: best_params.clone(),
            all_is_results: is_results,
            is_equity,
            oos_equity,
        });
        best_params_per_fold.push(best_params);
    }

    Ok(WalkForwardResult {
        folds: fold_results,
        best_params_per_fold,
    })
}

/// Run batch_lite on pre-aligned data (avoids reloading from disk).
fn run_on_pre_aligned(
    _runner: &SimpleBacktestRunner,
    strategies: &[StrategyPlan],
    config: &BacktestConfig,
    aligned: &bt_core::alignment::AlignedBars,
    universe_sorted: &[bt_kernel::SymbolId],
    max_parallelism: usize,
) -> Result<Vec<BatchResultLite>> {
    use bt_expr::DefaultExprEvaluator;
    use rayon::prelude::*;
    use rayon::ThreadPoolBuilder;

    let threads = if max_parallelism == 0 {
        10
    } else {
        max_parallelism
    };

    let pool = ThreadPoolBuilder::new()
        .num_threads(threads)
        .build()
        .map_err(|e| BtError::Strategy(format!("thread pool error: {e}")))?;

    let mut indexed = pool.install(|| {
        strategies
            .par_iter()
            .enumerate()
            .map(|(index, strategy)| {
                let evaluator = DefaultExprEvaluator::new();
                let r = SimpleBacktestRunner::with_evaluator(Box::new(evaluator));
                r.run_lite_on_aligned(strategy, config, aligned, universe_sorted)
                    .map(|result| (index, result))
            })
            .collect::<Vec<_>>()
    });

    indexed.sort_by_key(|entry| match entry {
        Ok((index, _)) => *index,
        Err(_) => usize::MAX,
    });

    let mut results = Vec::with_capacity(indexed.len());
    for entry in indexed {
        let (_, result) = entry?;
        results.push(result);
    }
    Ok(results)
}

/// Compute train/test time range pairs for each fold.
fn compute_fold_ranges(
    full_start: i64,
    full_end: i64,
    full_duration: i64,
    n_splits: usize,
    train_ratio: f64,
    method: &WalkForwardMethod,
) -> Result<Vec<(TimeRange, TimeRange)>> {
    let mut folds = Vec::with_capacity(n_splits);

    match method {
        WalkForwardMethod::Rolling => {
            // Each fold = train_ratio of segment for train, rest for test.
            let segment_size = full_duration / n_splits as i64;
            for i in 0..n_splits {
                let seg_start = full_start + i as i64 * segment_size;
                let seg_end = if i == n_splits - 1 {
                    full_end
                } else {
                    seg_start + segment_size
                };
                let split_point = seg_start + ((seg_end - seg_start) as f64 * train_ratio) as i64;

                let train = TimeRange::new(
                    Timestamp::from_nanos(seg_start),
                    Timestamp::from_nanos(split_point),
                )?;
                let test = TimeRange::new(
                    Timestamp::from_nanos(split_point),
                    Timestamp::from_nanos(seg_end),
                )?;
                folds.push((train, test));
            }
        }
        WalkForwardMethod::Anchored => {
            // Train grows from full_start; test is fixed-size segments.
            let test_total = (full_duration as f64 * (1.0 - train_ratio)) as i64;
            let test_size = test_total / n_splits as i64;

            for i in 0..n_splits {
                let test_start = full_end - test_total + i as i64 * test_size;
                let test_end = if i == n_splits - 1 {
                    full_end
                } else {
                    test_start + test_size
                };

                let train = TimeRange::new(
                    Timestamp::from_nanos(full_start),
                    Timestamp::from_nanos(test_start),
                )?;
                let test = TimeRange::new(
                    Timestamp::from_nanos(test_start),
                    Timestamp::from_nanos(test_end),
                )?;
                folds.push((train, test));
            }
        }
    }

    Ok(folds)
}
