//! Two-dimensional parameter sweep.
//!
//! Wraps `run_batch_lite` with a cartesian product of two parameter axes
//! and reshapes results into a 2D metric grid.

use std::collections::HashMap;

use bt_core::orchestrator::{BacktestConfig, StrategyPlan};
use bt_core::sweep::{expand_param_grid, strategy_with_params, ParamGrid};
use bt_core::{BacktestRunner, BatchResultLite, SimpleBacktestRunner};
use bt_data::DataStore;
use bt_expr::ScalarValue;
use bt_kernel::Result;
use serde::{Deserialize, Serialize};

use crate::extract_metric;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Sweep2dConfig {
    pub x_param: String,
    pub x_values: Vec<ScalarValue>,
    pub y_param: String,
    pub y_values: Vec<ScalarValue>,
    pub metric: String,
    pub max_parallelism: usize,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Sweep2dResult {
    pub x_param: String,
    pub x_values: Vec<ScalarValue>,
    pub y_param: String,
    pub y_values: Vec<ScalarValue>,
    /// `metric_grid[x_idx][y_idx]` — the metric value for the (x, y) combination.
    pub metric_grid: Vec<Vec<f64>>,
    /// Flat results in row-major order (x varies slowest).
    pub results: Vec<BatchResultLite>,
}

/// Run a 2D parameter sweep and return a metric grid.
pub fn run_sweep_2d(
    strategy: &StrategyPlan,
    sweep_config: &Sweep2dConfig,
    backtest_config: &BacktestConfig,
    data_store: &dyn DataStore,
) -> Result<Sweep2dResult> {
    let nx = sweep_config.x_values.len();
    let ny = sweep_config.y_values.len();

    // Build ParamGrid for cartesian expansion.
    let mut grid: ParamGrid = HashMap::new();
    grid.insert(sweep_config.x_param.clone(), sweep_config.x_values.clone());
    grid.insert(sweep_config.y_param.clone(), sweep_config.y_values.clone());

    let combos = expand_param_grid(&grid);

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

    // Run all combinations via batch_lite.
    let runner = SimpleBacktestRunner::new();
    let results = runner.run_batch_lite(
        &strategies,
        backtest_config,
        data_store,
        sweep_config.max_parallelism,
    )?;

    // expand_param_grid sorts keys alphabetically, so we need to figure out
    // which axis varies fastest. Keys are sorted, so the last key varies fastest.
    let mut sorted_keys: Vec<&str> = vec![&sweep_config.x_param, &sweep_config.y_param];
    sorted_keys.sort();

    let x_is_first = sorted_keys[0] == sweep_config.x_param;

    // Reshape flat results into 2D grid.
    let mut metric_grid = vec![vec![0.0_f64; ny]; nx];
    for (flat_idx, result) in results.iter().enumerate() {
        let (xi, yi) = if x_is_first {
            // x varies slowest (outer), y varies fastest (inner)
            (flat_idx / ny, flat_idx % ny)
        } else {
            // y varies slowest (outer), x varies fastest (inner)
            (flat_idx % nx, flat_idx / nx)
        };
        metric_grid[xi][yi] = extract_metric(&result.metrics, &sweep_config.metric);
    }

    Ok(Sweep2dResult {
        x_param: sweep_config.x_param.clone(),
        x_values: sweep_config.x_values.clone(),
        y_param: sweep_config.y_param.clone(),
        y_values: sweep_config.y_values.clone(),
        metric_grid,
        results,
    })
}
