//! Ensemble-level hyperopt adapter
//!
//! Adds N ensemble weight parameters (one per model) to a combined parameter space,
//! enabling joint optimization of model weights alongside per-model hyperparameters.
//!
//! ## Design
//!
//! The [`ParameterSpace`] trait uses static methods (`continuous_bounds()`, `param_names()`),
//! which means the dimensionality must be known at compile time. To support variable-size
//! ensemble configurations at runtime, this module provides:
//!
//! - [`EnsembleSpaceConfig`]: Runtime configuration specifying model names and per-model
//!   parameter dimensions. Stored in a thread-local so that static trait methods can
//!   access it.
//! - [`EnsembleParameterSpace`]: The parameter space struct implementing [`ParameterSpace`].
//!   Holds a flat continuous vector (per-model params concatenated + N weight values).
//!
//! ## Weight Normalization
//!
//! Raw weight values in `[0, 1]` are normalized via softmax-style division so they
//! sum to 1.0. If all raw weights are zero, equal weights are assigned.
//!
//! ## Usage
//!
//! ```rust,no_run
//! use ml::hyperopt::adapters::ensemble::{EnsembleSpaceConfig, EnsembleParameterSpace};
//! use ml::hyperopt::ParameterSpace;
//!
//! // Configure for a 3-model ensemble (DQN=11D, PPO=5D, TFT=6D)
//! let config = EnsembleSpaceConfig::new(
//!     vec!["dqn".into(), "ppo".into(), "tft".into()],
//!     vec![11, 5, 6],
//!     // Per-model bounds: 11 DQN bounds + 5 PPO bounds + 6 TFT bounds
//!     vec![
//!         // ... 22 bounds total from individual model ParameterSpaces
//!     ],
//!     // Per-model param names
//!     vec![
//!         // ... 22 names total
//!     ],
//! );
//!
//! // Install config so ParameterSpace static methods can read it
//! config.install();
//!
//! // Now EnsembleParameterSpace implements ParameterSpace
//! let bounds = EnsembleParameterSpace::continuous_bounds();
//! // bounds.len() == 22 (model params) + 3 (weights) = 25
//! ```

use std::cell::RefCell;

use crate::hyperopt::traits::ParameterSpace;
use crate::MLError;

// ---------------------------------------------------------------------------
// Runtime configuration
// ---------------------------------------------------------------------------

/// Runtime configuration for the ensemble parameter space.
///
/// Must be [`install()`](EnsembleSpaceConfig::install)ed before calling
/// `EnsembleParameterSpace` trait methods.
#[derive(Debug, Clone)]
pub struct EnsembleSpaceConfig {
    /// Human-readable model names (e.g., "dqn", "ppo", "tft").
    pub model_names: Vec<String>,
    /// Number of continuous parameters per model.
    pub model_param_dims: Vec<usize>,
    /// Concatenated per-model bounds `[(min, max), ...]`.
    /// Length must equal `model_param_dims.iter().sum()`.
    pub model_bounds: Vec<(f64, f64)>,
    /// Concatenated per-model parameter names.
    /// Length must equal `model_param_dims.iter().sum()`.
    pub model_param_names: Vec<String>,
}

thread_local! {
    static ENSEMBLE_CONFIG: RefCell<Option<EnsembleSpaceConfig>> = const { RefCell::new(None) };
}

impl EnsembleSpaceConfig {
    /// Create a new ensemble space configuration.
    ///
    /// # Arguments
    ///
    /// * `model_names` - Names of models in the ensemble.
    /// * `model_param_dims` - Number of hyperparameters per model.
    /// * `model_bounds` - Concatenated bounds for all model parameters.
    /// * `model_param_names` - Concatenated parameter names for all models.
    ///
    /// # Panics
    ///
    /// Panics if lengths are inconsistent.
    pub fn new(
        model_names: Vec<String>,
        model_param_dims: Vec<usize>,
        model_bounds: Vec<(f64, f64)>,
        model_param_names: Vec<String>,
    ) -> Self {
        let total_model_params: usize = model_param_dims.iter().sum();
        assert_eq!(
            model_names.len(),
            model_param_dims.len(),
            "model_names and model_param_dims must have the same length"
        );
        assert_eq!(
            model_bounds.len(),
            total_model_params,
            "model_bounds length must equal sum of model_param_dims"
        );
        assert_eq!(
            model_param_names.len(),
            total_model_params,
            "model_param_names length must equal sum of model_param_dims"
        );
        Self {
            model_names,
            model_param_dims,
            model_bounds,
            model_param_names,
        }
    }

    /// Install this configuration in the thread-local slot.
    ///
    /// Must be called before using `EnsembleParameterSpace` trait methods.
    pub fn install(&self) {
        ENSEMBLE_CONFIG.with(|cell| {
            *cell.borrow_mut() = Some(self.clone());
        });
    }

    /// Remove the installed configuration.
    pub fn uninstall() {
        ENSEMBLE_CONFIG.with(|cell| {
            *cell.borrow_mut() = None;
        });
    }

    /// Total dimension = sum(model_param_dims) + num_models (weight params).
    pub fn total_dim(&self) -> usize {
        let model_params: usize = self.model_param_dims.iter().sum();
        model_params + self.model_names.len()
    }

    /// Number of models in the ensemble.
    pub fn num_models(&self) -> usize {
        self.model_names.len()
    }
}

/// Read the installed config, returning an error if none is installed.
fn with_config<T>(f: impl FnOnce(&EnsembleSpaceConfig) -> T) -> Result<T, MLError> {
    ENSEMBLE_CONFIG.with(|cell| {
        let borrow = cell.borrow();
        match borrow.as_ref() {
            Some(cfg) => Ok(f(cfg)),
            None => Err(MLError::ConfigError("EnsembleSpaceConfig not installed. Call config.install() first."
                    .to_string())),
        }
    })
}

// ---------------------------------------------------------------------------
// Parameter space
// ---------------------------------------------------------------------------

/// Combined parameter space for ensemble optimization.
///
/// Holds a flat vector of continuous values:
/// `[model_0_params..., model_1_params..., ..., weight_0, weight_1, ...]`
///
/// The last N values (where N = number of models) are raw ensemble weights
/// in `[0, 1]`. Use [`extract_weights()`](Self::extract_weights) to get
/// normalized weights that sum to 1.
#[derive(Debug, Clone)]
pub struct EnsembleParameterSpace {
    /// Flat continuous parameter vector.
    pub values: Vec<f64>,
}

impl EnsembleParameterSpace {
    /// Extract normalized ensemble weights from the parameter vector.
    ///
    /// Weights are the last N values, normalized to sum to 1.
    /// If all raw weights are zero (or negative after clamping), returns equal weights.
    pub fn extract_weights(&self, num_models: usize) -> Vec<f64> {
        let weight_start = self.values.len().saturating_sub(num_models);
        let raw_weights: Vec<f64> = (0..num_models)
            .filter_map(|i| self.values.get(weight_start + i).copied())
            .map(|w| w.max(0.0))
            .collect();
        let sum: f64 = raw_weights.iter().sum();
        if sum > f64::EPSILON {
            raw_weights.iter().map(|w| w / sum).collect()
        } else {
            // Fallback: equal weights
            let equal = 1.0 / num_models.max(1) as f64;
            vec![equal; num_models]
        }
    }

    /// Extract per-model parameter slices from the combined vector.
    ///
    /// Returns a vector of slices, one per model. The slices are in the same
    /// order as the model names in the config.
    pub fn extract_model_params(&self, model_param_dims: &[usize]) -> Vec<Vec<f64>> {
        let mut result = Vec::with_capacity(model_param_dims.len());
        let mut offset = 0;
        for &dim in model_param_dims {
            let end = (offset + dim).min(self.values.len());
            let start = offset.min(end);
            result.push(self.values.get(start..end).unwrap_or_default().to_vec()); // cpu-side param split
            offset += dim;
        }
        result
    }
}

impl ParameterSpace for EnsembleParameterSpace {
    fn continuous_bounds() -> Vec<(f64, f64)> {
        // Read from thread-local config; if missing, return empty
        // (caller should have installed config first)
        with_config(|cfg| {
            let mut bounds = cfg.model_bounds.clone();
            // Append weight bounds: [0, 1] for each model
            for _ in 0..cfg.num_models() {
                bounds.push((0.0, 1.0));
            }
            bounds
        })
        .unwrap_or_default()
    }

    fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
        let expected_dim = with_config(|cfg| cfg.total_dim())?;
        if x.len() != expected_dim {
            return Err(MLError::ConfigError(format!(
                    "EnsembleParameterSpace: expected {} params, got {}",
                    expected_dim,
                    x.len()
                )));
        }
        Ok(Self {
            values: x.to_vec(), // cpu-side param clone
        })
    }

    fn to_continuous(&self) -> Vec<f64> {
        self.values.clone()
    }

    fn param_names() -> Vec<&'static str> {
        // Build names from config. Since the trait requires &'static str,
        // we leak the strings. This is acceptable because hyperopt configs
        // are created once per optimization run (not in a hot loop).
        with_config(|cfg| {
            let mut names: Vec<&'static str> = cfg
                .model_param_names
                .iter()
                .map(|s| -> &'static str { Box::leak(s.clone().into_boxed_str()) })
                .collect();
            for model_name in &cfg.model_names {
                let weight_name = format!("weight_{}", model_name);
                names.push(Box::leak(weight_name.into_boxed_str()));
            }
            names
        })
        .unwrap_or_default()
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    /// Helper: create a simple 2-model config for testing.
    fn test_config() -> EnsembleSpaceConfig {
        EnsembleSpaceConfig::new(
            vec!["model_a".into(), "model_b".into()],
            vec![2, 3], // model_a has 2 params, model_b has 3 params
            vec![
                (0.0, 1.0),
                (0.0, 10.0),
                (-1.0, 1.0),
                (-1.0, 1.0),
                (0.0, 100.0),
            ],
            vec![
                "a_lr".into(),
                "a_batch".into(),
                "b_x".into(),
                "b_y".into(),
                "b_z".into(),
            ],
        )
    }

    #[test]
    fn test_ensemble_parameter_space_dimensions() {
        let config = test_config();
        config.install();

        // 2 + 3 model params + 2 weights = 7 total
        assert_eq!(config.total_dim(), 7);

        let bounds = EnsembleParameterSpace::continuous_bounds();
        assert_eq!(bounds.len(), 7);

        // First 5 bounds are model params
        assert_eq!(bounds.first().copied(), Some((0.0, 1.0))); // a_lr
        assert_eq!(bounds.get(4).copied(), Some((0.0, 100.0))); // b_z

        // Last 2 bounds are weight params [0, 1]
        assert_eq!(bounds.get(5).copied(), Some((0.0, 1.0))); // weight_model_a
        assert_eq!(bounds.get(6).copied(), Some((0.0, 1.0))); // weight_model_b

        let names = EnsembleParameterSpace::param_names();
        assert_eq!(names.len(), 7);
        assert_eq!(names.first().copied(), Some("a_lr"));
        assert_eq!(names.get(4).copied(), Some("b_z"));
        assert_eq!(names.get(5).copied(), Some("weight_model_a"));
        assert_eq!(names.get(6).copied(), Some("weight_model_b"));

        EnsembleSpaceConfig::uninstall();
    }

    #[test]
    fn test_extract_weights_normalizes() {
        let config = test_config();
        config.install();

        // 7 params: [a_lr, a_batch, b_x, b_y, b_z, weight_a, weight_b]
        let params = EnsembleParameterSpace::from_continuous(&[
            0.5, 5.0, 0.0, 0.0, 50.0, // model params
            0.3, 0.7, // raw weights
        ])
        .unwrap_or_else(|e| panic!("from_continuous failed: {}", e));

        let weights = params.extract_weights(2);
        assert_eq!(weights.len(), 2);

        // 0.3 / (0.3 + 0.7) = 0.3
        assert!((weights.first().copied().unwrap_or(0.0) - 0.3).abs() < 1e-10);
        // 0.7 / (0.3 + 0.7) = 0.7
        assert!((weights.get(1).copied().unwrap_or(0.0) - 0.7).abs() < 1e-10);

        // Verify weights sum to 1
        let sum: f64 = weights.iter().sum();
        assert!(
            (sum - 1.0).abs() < 1e-10,
            "Weights should sum to 1.0, got {}",
            sum
        );

        EnsembleSpaceConfig::uninstall();
    }

    #[test]
    fn test_extract_weights_handles_zeros() {
        let config = test_config();
        config.install();

        let params = EnsembleParameterSpace::from_continuous(&[
            0.5, 5.0, 0.0, 0.0, 50.0, // model params
            0.0, 0.0, // all-zero weights
        ])
        .unwrap_or_else(|e| panic!("from_continuous failed: {}", e));

        let weights = params.extract_weights(2);
        assert_eq!(weights.len(), 2);

        // Should fall back to equal weights: 0.5, 0.5
        assert!(
            (weights.first().copied().unwrap_or(0.0) - 0.5).abs() < 1e-10,
            "Expected 0.5 for zero-weight fallback, got {:?}",
            weights.first()
        );
        assert!(
            (weights.get(1).copied().unwrap_or(0.0) - 0.5).abs() < 1e-10,
            "Expected 0.5 for zero-weight fallback, got {:?}",
            weights.get(1)
        );

        EnsembleSpaceConfig::uninstall();
    }

    #[test]
    fn test_roundtrip_continuous() {
        let config = test_config();
        config.install();

        let original = vec![0.5, 5.0, 0.0, -0.5, 50.0, 0.3, 0.7];
        let params = EnsembleParameterSpace::from_continuous(&original)
            .unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
        let recovered = params.to_continuous();

        assert_eq!(original.len(), recovered.len());
        for (a, b) in original.iter().zip(recovered.iter()) {
            assert!(
                (a - b).abs() < 1e-10,
                "Round-trip mismatch: {} vs {}",
                a,
                b
            );
        }

        EnsembleSpaceConfig::uninstall();
    }

    #[test]
    fn test_wrong_dimension_returns_error() {
        let config = test_config();
        config.install();

        let result = EnsembleParameterSpace::from_continuous(&[1.0, 2.0]);
        assert!(result.is_err(), "Should reject wrong dimension");

        EnsembleSpaceConfig::uninstall();
    }

    #[test]
    fn test_extract_model_params() {
        let config = test_config();
        config.install();

        let params = EnsembleParameterSpace::from_continuous(&[
            0.5, 5.0, // model_a params (dim=2)
            0.1, -0.3, 75.0, // model_b params (dim=3)
            0.4, 0.6, // weights
        ])
        .unwrap_or_else(|e| panic!("from_continuous failed: {}", e));

        let model_params = params.extract_model_params(&config.model_param_dims);
        assert_eq!(model_params.len(), 2);

        // model_a: [0.5, 5.0]
        assert_eq!(model_params.first().map(|v| v.len()), Some(2));
        assert!(
            (model_params
                .first()
                .and_then(|v| v.first().copied())
                .unwrap_or(0.0)
                - 0.5)
                .abs()
                < 1e-10
        );

        // model_b: [0.1, -0.3, 75.0]
        assert_eq!(model_params.get(1).map(|v| v.len()), Some(3));
        assert!(
            (model_params
                .get(1)
                .and_then(|v| v.get(2).copied())
                .unwrap_or(0.0)
                - 75.0)
                .abs()
                < 1e-10
        );

        EnsembleSpaceConfig::uninstall();
    }

    #[test]
    fn test_no_config_installed_returns_empty_or_error() {
        // Ensure no config is installed
        EnsembleSpaceConfig::uninstall();

        // continuous_bounds returns empty when no config
        let bounds = EnsembleParameterSpace::continuous_bounds();
        assert!(bounds.is_empty());

        // from_continuous returns error when no config
        let result = EnsembleParameterSpace::from_continuous(&[1.0]);
        assert!(result.is_err());

        // param_names returns empty when no config
        let names = EnsembleParameterSpace::param_names();
        assert!(names.is_empty());
    }

    #[test]
    fn test_single_model_ensemble() {
        let config = EnsembleSpaceConfig::new(
            vec!["only_model".into()],
            vec![1],
            vec![(0.0, 1.0)],
            vec!["lr".into()],
        );
        config.install();

        assert_eq!(config.total_dim(), 2); // 1 param + 1 weight

        let params = EnsembleParameterSpace::from_continuous(&[0.5, 0.8])
            .unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
        let weights = params.extract_weights(1);
        assert_eq!(weights.len(), 1);
        // Single model: weight normalizes to 1.0
        assert!(
            (weights.first().copied().unwrap_or(0.0) - 1.0).abs() < 1e-10,
            "Single model weight should be 1.0"
        );

        EnsembleSpaceConfig::uninstall();
    }
}
