//! Tree-Parzen Estimator (TPE) sampler implementation.
//!
//! TPE is a Bayesian optimization algorithm that models the objective function
//! using two probability distributions: one for promising (good) parameter values
//! and one for unpromising (bad) parameter values.
//!
//! # Gamma Strategies
//!
//! The gamma parameter controls what fraction of trials are considered "good".
//! This module provides several built-in strategies via the [`GammaStrategy`] trait:
//!
//! - [`FixedGamma`]: Constant gamma value (default: 0.25)
//! - [`LinearGamma`]: Linear interpolation between min and max based on trial count
//! - [`SqrtGamma`]: Gamma decreases as 1/√n (similar to Optuna)
//! - [`HyperoptGamma`]: Hyperopt-style adaptive gamma
//!
//! You can also implement your own strategy by implementing the [`GammaStrategy`] trait.
//!
//! # Examples
//!
//! Using a built-in gamma strategy:
//!
//! ```
//! use optimizer::sampler::tpe::{SqrtGamma, TpeSampler};
//!
//! let sampler = TpeSampler::builder()
//!     .gamma_strategy(SqrtGamma::default())
//!     .build()
//!     .unwrap();
//! ```
//!
//! Implementing a custom gamma strategy:
//!
//! ```
//! use optimizer::sampler::tpe::{GammaStrategy, TpeSampler};
//!
//! #[derive(Debug, Clone)]
//! struct MyGamma {
//!     base: f64,
//! }
//!
//! impl GammaStrategy for MyGamma {
//!     fn gamma(&self, n_trials: usize) -> f64 {
//!         (self.base + 0.01 * n_trials as f64).min(0.5)
//!     }
//!
//!     fn clone_box(&self) -> Box<dyn GammaStrategy> {
//!         Box::new(self.clone())
//!     }
//! }
//!
//! let sampler = TpeSampler::builder()
//!     .gamma_strategy(MyGamma { base: 0.1 })
//!     .build()
//!     .unwrap();
//! ```

use core::fmt::Debug;
use std::sync::Arc;

use parking_lot::Mutex;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};

use crate::distribution::Distribution;
use crate::error::{Error, Result};
use crate::kde::KernelDensityEstimator;
use crate::param::ParamValue;
use crate::sampler::{CompletedTrial, Sampler};

// ============================================================================
// Gamma Strategy Trait and Implementations
// ============================================================================

/// A strategy for computing the gamma quantile in TPE.
///
/// The gamma value determines what fraction of trials are considered "good"
/// when splitting the trial history. Different strategies can adapt this
/// fraction based on the number of completed trials.
///
/// # Implementation Notes
///
/// - The returned gamma must be in the range (0.0, 1.0)
/// - Implementations should be deterministic for reproducibility
/// - The `clone_box` method enables trait object cloning
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::GammaStrategy;
///
/// #[derive(Debug, Clone)]
/// struct ConstantGamma(f64);
///
/// impl GammaStrategy for ConstantGamma {
///     fn gamma(&self, _n_trials: usize) -> f64 {
///         self.0
///     }
///
///     fn clone_box(&self) -> Box<dyn GammaStrategy> {
///         Box::new(self.clone())
///     }
/// }
/// ```
pub trait GammaStrategy: Send + Sync + Debug {
    /// Computes the gamma quantile based on the number of completed trials.
    ///
    /// # Arguments
    ///
    /// * `n_trials` - The number of completed trials in the history.
    ///
    /// # Returns
    ///
    /// A gamma value in the range (0.0, 1.0). Values outside this range
    /// will be clamped by the sampler.
    fn gamma(&self, n_trials: usize) -> f64;

    /// Creates a boxed clone of this strategy.
    ///
    /// This method enables cloning of trait objects, which is necessary
    /// for the builder pattern and sampler configuration.
    fn clone_box(&self) -> Box<dyn GammaStrategy>;
}

impl Clone for Box<dyn GammaStrategy> {
    fn clone(&self) -> Self {
        self.clone_box()
    }
}

/// A fixed gamma strategy that returns a constant value.
///
/// This is the simplest strategy and the default behavior of TPE.
/// The gamma value remains constant regardless of the number of trials.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{FixedGamma, TpeSampler};
///
/// // Use 15% of trials as "good"
/// let sampler = TpeSampler::builder()
///     .gamma_strategy(FixedGamma::new(0.15).unwrap())
///     .build()
///     .unwrap();
/// ```
#[derive(Debug, Clone, Copy)]
pub struct FixedGamma {
    gamma: f64,
}

impl FixedGamma {
    /// Creates a new fixed gamma strategy.
    ///
    /// # Arguments
    ///
    /// * `gamma` - The constant gamma value to use.
    ///
    /// # Errors
    ///
    /// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0).
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::FixedGamma;
    ///
    /// let strategy = FixedGamma::new(0.25).unwrap();
    /// assert!((strategy.value() - 0.25).abs() < f64::EPSILON);
    /// ```
    pub fn new(gamma: f64) -> Result<Self> {
        if gamma <= 0.0 || gamma >= 1.0 {
            return Err(Error::InvalidGamma(gamma));
        }
        Ok(Self { gamma })
    }

    /// Returns the fixed gamma value.
    #[must_use]
    pub fn value(&self) -> f64 {
        self.gamma
    }
}

impl Default for FixedGamma {
    /// Creates a fixed gamma strategy with the default value of 0.25.
    fn default() -> Self {
        Self { gamma: 0.25 }
    }
}

impl GammaStrategy for FixedGamma {
    fn gamma(&self, _n_trials: usize) -> f64 {
        self.gamma
    }

    fn clone_box(&self) -> Box<dyn GammaStrategy> {
        Box::new(*self)
    }
}

/// A linear gamma strategy that interpolates between min and max values.
///
/// The gamma value increases linearly from `gamma_min` to `gamma_max` as the
/// number of trials grows from 0 to `n_trials_max`. Beyond `n_trials_max`,
/// gamma remains at `gamma_max`.
///
/// This strategy is useful when you want to be more explorative early on
/// (smaller gamma = fewer "good" trials) and more exploitative later
/// (larger gamma = more "good" trials).
///
/// # Formula
///
/// ```text
/// gamma = gamma_min + (gamma_max - gamma_min) * min(n_trials / n_trials_max, 1.0)
/// ```
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, LinearGamma, TpeSampler};
///
/// let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap();
///
/// // At 0 trials: gamma = 0.1
/// assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON);
///
/// // At 50 trials: gamma = 0.25 (midpoint)
/// assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON);
///
/// // At 100+ trials: gamma = 0.4
/// assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON);
/// assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON);
/// ```
#[derive(Debug, Clone, Copy)]
pub struct LinearGamma {
    gamma_min: f64,
    gamma_max: f64,
    n_trials_max: usize,
}

impl LinearGamma {
    /// Creates a new linear gamma strategy.
    ///
    /// # Arguments
    ///
    /// * `gamma_min` - The minimum gamma value (at 0 trials).
    /// * `gamma_max` - The maximum gamma value (at `n_trials_max` trials).
    /// * `n_trials_max` - The number of trials at which gamma reaches its maximum.
    ///
    /// # Errors
    ///
    /// Returns `Error::InvalidGamma` if:
    /// - `gamma_min` is not in (0.0, 1.0)
    /// - `gamma_max` is not in (0.0, 1.0)
    /// - `gamma_min > gamma_max`
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::LinearGamma;
    ///
    /// // Gamma goes from 0.1 to 0.3 over 50 trials
    /// let strategy = LinearGamma::new(0.1, 0.3, 50).unwrap();
    /// ```
    pub fn new(gamma_min: f64, gamma_max: f64, n_trials_max: usize) -> Result<Self> {
        if gamma_min <= 0.0 || gamma_min >= 1.0 {
            return Err(Error::InvalidGamma(gamma_min));
        }
        if gamma_max <= 0.0 || gamma_max >= 1.0 {
            return Err(Error::InvalidGamma(gamma_max));
        }
        if gamma_min > gamma_max {
            return Err(Error::InvalidGamma(gamma_min));
        }
        Ok(Self {
            gamma_min,
            gamma_max,
            n_trials_max,
        })
    }

    /// Returns the minimum gamma value.
    #[must_use]
    pub fn gamma_min(&self) -> f64 {
        self.gamma_min
    }

    /// Returns the maximum gamma value.
    #[must_use]
    pub fn gamma_max(&self) -> f64 {
        self.gamma_max
    }

    /// Returns the number of trials at which gamma reaches its maximum.
    #[must_use]
    pub fn n_trials_max(&self) -> usize {
        self.n_trials_max
    }
}

impl Default for LinearGamma {
    /// Creates a linear gamma strategy with default values:
    /// - `gamma_min`: 0.10
    /// - `gamma_max`: 0.25
    /// - `n_trials_max`: 100
    fn default() -> Self {
        Self {
            gamma_min: 0.10,
            gamma_max: 0.25,
            n_trials_max: 100,
        }
    }
}

impl GammaStrategy for LinearGamma {
    #[allow(clippy::cast_precision_loss)]
    fn gamma(&self, n_trials: usize) -> f64 {
        if self.n_trials_max == 0 {
            return self.gamma_max;
        }
        let t = (n_trials as f64 / self.n_trials_max as f64).min(1.0);
        self.gamma_min + (self.gamma_max - self.gamma_min) * t
    }

    fn clone_box(&self) -> Box<dyn GammaStrategy> {
        Box::new(*self)
    }
}

/// A square root gamma strategy inspired by Optuna's default behavior.
///
/// The gamma value is computed based on the inverse square root of the number
/// of trials, providing a balance between exploration and exploitation that
/// naturally adapts as more data becomes available.
///
/// # Formula
///
/// ```text
/// n_good = max(1, floor(gamma_factor / sqrt(n_trials)))
/// gamma = min(gamma_max, n_good / n_trials)
/// ```
///
/// When `n_trials` is 0, returns `gamma_max`.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, SqrtGamma, TpeSampler};
///
/// let strategy = SqrtGamma::default();
///
/// // Gamma decreases as trials increase
/// let g10 = strategy.gamma(10);
/// let g100 = strategy.gamma(100);
/// assert!(g10 > g100, "Gamma should decrease with more trials");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct SqrtGamma {
    gamma_factor: f64,
    gamma_max: f64,
}

impl SqrtGamma {
    /// Creates a new square root gamma strategy.
    ///
    /// # Arguments
    ///
    /// * `gamma_factor` - The factor controlling how quickly gamma decreases.
    ///   Higher values mean more "good" trials at any given point.
    /// * `gamma_max` - The maximum gamma value (used when `n_trials` is small).
    ///
    /// # Errors
    ///
    /// Returns `Error::InvalidGamma` if:
    /// - `gamma_factor` is not positive
    /// - `gamma_max` is not in (0.0, 1.0)
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::SqrtGamma;
    ///
    /// let strategy = SqrtGamma::new(1.0, 0.25).unwrap();
    /// ```
    pub fn new(gamma_factor: f64, gamma_max: f64) -> Result<Self> {
        if gamma_factor <= 0.0 {
            return Err(Error::InvalidGamma(gamma_factor));
        }
        if gamma_max <= 0.0 || gamma_max >= 1.0 {
            return Err(Error::InvalidGamma(gamma_max));
        }
        Ok(Self {
            gamma_factor,
            gamma_max,
        })
    }

    /// Returns the gamma factor.
    #[must_use]
    pub fn gamma_factor(&self) -> f64 {
        self.gamma_factor
    }

    /// Returns the maximum gamma value.
    #[must_use]
    pub fn gamma_max(&self) -> f64 {
        self.gamma_max
    }
}

impl Default for SqrtGamma {
    /// Creates a square root gamma strategy with default values:
    /// - `gamma_factor`: 1.0
    /// - `gamma_max`: 0.25
    fn default() -> Self {
        Self {
            gamma_factor: 1.0,
            gamma_max: 0.25,
        }
    }
}

impl GammaStrategy for SqrtGamma {
    #[allow(clippy::cast_precision_loss)]
    fn gamma(&self, n_trials: usize) -> f64 {
        if n_trials == 0 {
            return self.gamma_max;
        }
        let n_good = (self.gamma_factor / (n_trials as f64).sqrt()).max(1.0);
        (n_good / n_trials as f64).min(self.gamma_max)
    }

    fn clone_box(&self) -> Box<dyn GammaStrategy> {
        Box::new(*self)
    }
}

/// A Hyperopt-style gamma strategy.
///
/// This strategy computes gamma as `min(gamma_max, (gamma_base + 1) / n_trials)`,
/// which is inspired by the original Hyperopt TPE implementation.
///
/// # Formula
///
/// ```text
/// gamma = min(gamma_max, (gamma_base + 1) / n_trials)
/// ```
///
/// When `n_trials` is 0, returns `gamma_max`.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, HyperoptGamma};
///
/// // With gamma_base=24 and gamma_max=0.5:
/// // - At n=25: gamma = min(0.5, 25/25) = 0.5 (capped)
/// // - At n=100: gamma = min(0.5, 25/100) = 0.25
/// let strategy = HyperoptGamma::new(24.0, 0.5).unwrap();
///
/// // Early trials have higher gamma
/// let g50 = strategy.gamma(50);
/// let g200 = strategy.gamma(200);
/// assert!(g50 > g200, "Gamma should decrease with more trials");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct HyperoptGamma {
    gamma_base: f64,
    gamma_max: f64,
}

impl HyperoptGamma {
    /// Creates a new Hyperopt-style gamma strategy.
    ///
    /// # Arguments
    ///
    /// * `gamma_base` - The base value added to 1 in the numerator.
    /// * `gamma_max` - The maximum gamma value.
    ///
    /// # Errors
    ///
    /// Returns `Error::InvalidGamma` if:
    /// - `gamma_base` is negative
    /// - `gamma_max` is not in (0.0, 1.0)
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::HyperoptGamma;
    ///
    /// let strategy = HyperoptGamma::new(24.0, 0.25).unwrap();
    /// ```
    pub fn new(gamma_base: f64, gamma_max: f64) -> Result<Self> {
        if gamma_base < 0.0 {
            return Err(Error::InvalidGamma(gamma_base));
        }
        if gamma_max <= 0.0 || gamma_max >= 1.0 {
            return Err(Error::InvalidGamma(gamma_max));
        }
        Ok(Self {
            gamma_base,
            gamma_max,
        })
    }

    /// Returns the gamma base value.
    #[must_use]
    pub fn gamma_base(&self) -> f64 {
        self.gamma_base
    }

    /// Returns the maximum gamma value.
    #[must_use]
    pub fn gamma_max(&self) -> f64 {
        self.gamma_max
    }
}

impl Default for HyperoptGamma {
    /// Creates a Hyperopt-style gamma strategy with default values:
    /// - `gamma_base`: 24.0
    /// - `gamma_max`: 0.25
    fn default() -> Self {
        Self {
            gamma_base: 24.0,
            gamma_max: 0.25,
        }
    }
}

impl GammaStrategy for HyperoptGamma {
    #[allow(clippy::cast_precision_loss)]
    fn gamma(&self, n_trials: usize) -> f64 {
        if n_trials == 0 {
            return self.gamma_max;
        }
        ((self.gamma_base + 1.0) / n_trials as f64).min(self.gamma_max)
    }

    fn clone_box(&self) -> Box<dyn GammaStrategy> {
        Box::new(*self)
    }
}

// ============================================================================
// TPE Sampler
// ============================================================================

/// A Tree-Parzen Estimator (TPE) sampler for Bayesian optimization.
///
/// TPE works by splitting completed trials into two groups based on their
/// objective values: good trials (below the gamma quantile) and bad trials
/// (above the gamma quantile). It then fits kernel density estimators (KDE)
/// to each group and samples new points that maximize the ratio l(x)/g(x),
/// where l(x) is the density of good trials and g(x) is the density of bad trials.
///
/// During the startup phase (when fewer than `n_startup_trials` are completed),
/// TPE falls back to random sampling to gather initial data.
///
/// # Gamma Strategies
///
/// The gamma quantile can be configured using different strategies via the
/// [`GammaStrategy`] trait. Built-in strategies include:
///
/// - [`FixedGamma`]: Constant gamma (default: 0.25)
/// - [`LinearGamma`]: Linear interpolation based on trial count
/// - [`SqrtGamma`]: Inverse square root scaling (Optuna-style)
/// - [`HyperoptGamma`]: Hyperopt-style adaptive gamma
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::TpeSampler;
///
/// // Create with default settings (FixedGamma at 0.25)
/// let sampler = TpeSampler::new();
///
/// // Create with custom settings using the builder
/// let sampler = TpeSampler::builder()
///     .gamma(0.15)  // Shorthand for FixedGamma::new(0.15)
///     .n_startup_trials(20)
///     .n_ei_candidates(32)
///     .seed(42)
///     .build()
///     .unwrap();
/// ```
///
/// Using a different gamma strategy:
///
/// ```
/// use optimizer::sampler::tpe::{SqrtGamma, TpeSampler};
///
/// let sampler = TpeSampler::builder()
///     .gamma_strategy(SqrtGamma::default())
///     .build()
///     .unwrap();
/// ```
pub struct TpeSampler {
    /// Strategy for computing the gamma quantile.
    gamma_strategy: Arc<dyn GammaStrategy>,
    /// Number of trials before TPE kicks in (uses random sampling before this).
    n_startup_trials: usize,
    /// Number of candidate samples to evaluate when selecting the next point.
    n_ei_candidates: usize,
    /// Optional fixed bandwidth for KDE. If None, uses Scott's rule.
    kde_bandwidth: Option<f64>,
    /// Thread-safe RNG for sampling.
    rng: Mutex<StdRng>,
}

impl TpeSampler {
    /// Creates a new TPE sampler with default settings.
    ///
    /// Default settings:
    /// - gamma strategy: [`FixedGamma`] with gamma = 0.25
    /// - `n_startup_trials`: 10 (random sampling for first 10 trials)
    /// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample)
    /// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth)
    #[must_use]
    pub fn new() -> Self {
        Self {
            gamma_strategy: Arc::new(FixedGamma::default()),
            n_startup_trials: 10,
            n_ei_candidates: 24,
            kde_bandwidth: None,
            rng: Mutex::new(StdRng::from_os_rng()),
        }
    }

    /// Creates a builder for configuring a TPE sampler.
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::TpeSampler;
    ///
    /// let sampler = TpeSampler::builder()
    ///     .gamma(0.15)
    ///     .n_startup_trials(20)
    ///     .n_ei_candidates(32)
    ///     .seed(42)
    ///     .build()
    ///     .unwrap();
    /// ```
    #[must_use]
    pub fn builder() -> TpeSamplerBuilder {
        TpeSamplerBuilder::new()
    }

    /// Creates a new TPE sampler with custom configuration.
    ///
    /// This method uses a fixed gamma value. For more advanced gamma strategies,
    /// use [`TpeSampler::with_strategy`] or the builder pattern with
    /// [`TpeSamplerBuilder::gamma_strategy`].
    ///
    /// # Arguments
    ///
    /// * `gamma` - Fraction of trials to consider "good" (0.0 to 1.0).
    /// * `n_startup_trials` - Number of random trials before TPE sampling.
    /// * `n_ei_candidates` - Number of candidates to evaluate per sample.
    /// * `kde_bandwidth` - Optional fixed bandwidth for KDE. If None, uses Scott's rule.
    /// * `seed` - Optional seed for reproducibility.
    ///
    /// # Errors
    ///
    /// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0).
    /// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
    pub fn with_config(
        gamma: f64,
        n_startup_trials: usize,
        n_ei_candidates: usize,
        kde_bandwidth: Option<f64>,
        seed: Option<u64>,
    ) -> Result<Self> {
        let gamma_strategy = FixedGamma::new(gamma)?;
        Self::with_strategy(
            gamma_strategy,
            n_startup_trials,
            n_ei_candidates,
            kde_bandwidth,
            seed,
        )
    }

    /// Creates a new TPE sampler with a custom gamma strategy.
    ///
    /// # Arguments
    ///
    /// * `gamma_strategy` - The strategy for computing the gamma quantile.
    /// * `n_startup_trials` - Number of random trials before TPE sampling.
    /// * `n_ei_candidates` - Number of candidates to evaluate per sample.
    /// * `kde_bandwidth` - Optional fixed bandwidth for KDE. If None, uses Scott's rule.
    /// * `seed` - Optional seed for reproducibility.
    ///
    /// # Errors
    ///
    /// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::{SqrtGamma, TpeSampler};
    ///
    /// let sampler = TpeSampler::with_strategy(
    ///     SqrtGamma::default(),
    ///     10,       // n_startup_trials
    ///     24,       // n_ei_candidates
    ///     None,     // kde_bandwidth
    ///     Some(42), // seed
    /// )
    /// .unwrap();
    /// ```
    pub fn with_strategy<G: GammaStrategy + 'static>(
        gamma_strategy: G,
        n_startup_trials: usize,
        n_ei_candidates: usize,
        kde_bandwidth: Option<f64>,
        seed: Option<u64>,
    ) -> Result<Self> {
        if let Some(bw) = kde_bandwidth
            && bw <= 0.0
        {
            return Err(Error::InvalidBandwidth(bw));
        }

        let rng = match seed {
            Some(s) => StdRng::seed_from_u64(s),
            None => StdRng::from_os_rng(),
        };

        Ok(Self {
            gamma_strategy: Arc::new(gamma_strategy),
            n_startup_trials,
            n_ei_candidates,
            kde_bandwidth,
            rng: Mutex::new(rng),
        })
    }

    /// Returns the gamma strategy used by this sampler.
    #[must_use]
    pub fn gamma_strategy(&self) -> &dyn GammaStrategy {
        self.gamma_strategy.as_ref()
    }

    /// Splits trials into good and bad groups based on the gamma quantile.
    ///
    /// The gamma value is computed dynamically using the configured [`GammaStrategy`].
    ///
    /// Returns (`good_trials`, `bad_trials`) where `good_trials` contains trials
    /// with values below the gamma quantile (for minimization).
    #[allow(
        clippy::cast_precision_loss,
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss
    )]
    fn split_trials<'a>(
        &self,
        history: &'a [CompletedTrial],
    ) -> (Vec<&'a CompletedTrial>, Vec<&'a CompletedTrial>) {
        if history.is_empty() {
            return (vec![], vec![]);
        }

        // Sort trials by value (ascending for minimization)
        let mut sorted_indices: Vec<usize> = (0..history.len()).collect();
        sorted_indices.sort_by(|&a, &b| {
            history[a]
                .value
                .partial_cmp(&history[b].value)
                .unwrap_or(core::cmp::Ordering::Equal)
        });

        // Compute gamma using the strategy and clamp to valid range
        let gamma = self
            .gamma_strategy
            .gamma(history.len())
            .clamp(f64::EPSILON, 1.0 - f64::EPSILON);

        // Calculate the split point (gamma quantile)
        // Ensure at least 1 trial in each group if possible
        let n_good = ((history.len() as f64 * gamma).ceil() as usize)
            .max(1)
            .min(history.len() - 1);

        let good: Vec<_> = sorted_indices[..n_good]
            .iter()
            .map(|&i| &history[i])
            .collect();
        let bad: Vec<_> = sorted_indices[n_good..]
            .iter()
            .map(|&i| &history[i])
            .collect();

        (good, bad)
    }

    /// Samples uniformly from a distribution (used during startup phase).
    #[allow(
        clippy::cast_possible_truncation,
        clippy::cast_precision_loss,
        clippy::unused_self
    )]
    fn sample_uniform(&self, distribution: &Distribution, rng: &mut StdRng) -> ParamValue {
        match distribution {
            Distribution::Float(d) => {
                let value = if d.log_scale {
                    let log_low = d.low.ln();
                    let log_high = d.high.ln();
                    rng.random_range(log_low..=log_high).exp()
                } else if let Some(step) = d.step {
                    let n_steps = ((d.high - d.low) / step).floor() as i64;
                    let k = rng.random_range(0..=n_steps);
                    d.low + (k as f64) * step
                } else {
                    rng.random_range(d.low..=d.high)
                };
                ParamValue::Float(value)
            }
            Distribution::Int(d) => {
                let value = if d.log_scale {
                    let log_low = (d.low as f64).ln();
                    let log_high = (d.high as f64).ln();
                    let raw = rng.random_range(log_low..=log_high).exp().round() as i64;
                    raw.clamp(d.low, d.high)
                } else if let Some(step) = d.step {
                    let n_steps = (d.high - d.low) / step;
                    let k = rng.random_range(0..=n_steps);
                    d.low + k * step
                } else {
                    rng.random_range(d.low..=d.high)
                };
                ParamValue::Int(value)
            }
            Distribution::Categorical(d) => {
                ParamValue::Categorical(rng.random_range(0..d.n_choices))
            }
        }
    }

    /// Samples using TPE for float distributions.
    #[allow(clippy::too_many_arguments)]
    fn sample_tpe_float(
        &self,
        low: f64,
        high: f64,
        log_scale: bool,
        step: Option<f64>,
        good_values: Vec<f64>,
        bad_values: Vec<f64>,
        rng: &mut StdRng,
    ) -> f64 {
        // Transform to internal space (log space if needed)
        let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
            let i_low = low.ln();
            let i_high = high.ln();
            let g: Vec<f64> = good_values.iter().map(|&v| v.ln()).collect();
            let b: Vec<f64> = bad_values.iter().map(|&v| v.ln()).collect();
            (i_low, i_high, g, b)
        } else {
            (low, high, good_values, bad_values)
        };

        // Fit KDEs to good and bad groups
        let l_kde = match self.kde_bandwidth {
            Some(bw) => KernelDensityEstimator::with_bandwidth(good_internal, bw),
            None => KernelDensityEstimator::new(good_internal),
        };
        let g_kde = match self.kde_bandwidth {
            Some(bw) => KernelDensityEstimator::with_bandwidth(bad_internal, bw),
            None => KernelDensityEstimator::new(bad_internal),
        };

        // If KDE construction fails, fall back to uniform sampling
        let (Ok(l_kde), Ok(g_kde)) = (l_kde, g_kde) else {
            return rng.random_range(low..=high);
        };

        // Generate candidates from l(x) and select the one with best l(x)/g(x) ratio
        let mut best_candidate = internal_low;
        let mut best_ratio = f64::NEG_INFINITY;

        for _ in 0..self.n_ei_candidates {
            let candidate = l_kde.sample(rng);

            // Clamp to bounds
            let candidate = candidate.clamp(internal_low, internal_high);

            let l_density = l_kde.pdf(candidate);
            let g_density = g_kde.pdf(candidate);

            // Compute l(x)/g(x) ratio, handling zero density
            let ratio = if g_density < f64::EPSILON {
                if l_density > f64::EPSILON {
                    f64::INFINITY
                } else {
                    0.0
                }
            } else {
                l_density / g_density
            };

            if ratio > best_ratio {
                best_ratio = ratio;
                best_candidate = candidate;
            }
        }

        // Transform back from internal space
        let mut value = if log_scale {
            best_candidate.exp()
        } else {
            best_candidate
        };

        // Apply step constraint if present
        if let Some(step) = step {
            let k = ((value - low) / step).round();
            value = low + k * step;
        }

        // Ensure value is within bounds
        value.clamp(low, high)
    }

    /// Samples using TPE for integer distributions.
    #[allow(
        clippy::too_many_arguments,
        clippy::cast_precision_loss,
        clippy::cast_possible_truncation
    )]
    fn sample_tpe_int(
        &self,
        low: i64,
        high: i64,
        log_scale: bool,
        step: Option<i64>,
        good_values: &[i64],
        bad_values: &[i64],
        rng: &mut StdRng,
    ) -> i64 {
        // Convert to floats for KDE
        let good_floats: Vec<f64> = good_values.iter().map(|&v| v as f64).collect();
        let bad_floats: Vec<f64> = bad_values.iter().map(|&v| v as f64).collect();

        // Use float TPE sampling
        let float_value = self.sample_tpe_float(
            low as f64,
            high as f64,
            log_scale,
            step.map(|s| s as f64),
            good_floats,
            bad_floats,
            rng,
        );

        // Round to nearest integer
        let int_value = float_value.round() as i64;

        // Apply step constraint if present
        let int_value = if let Some(step) = step {
            let k = ((int_value - low) as f64 / step as f64).round() as i64;
            low + k * step
        } else {
            int_value
        };

        // Ensure value is within bounds
        int_value.clamp(low, high)
    }

    /// Samples using TPE for categorical distributions.
    #[allow(clippy::cast_precision_loss, clippy::unused_self)]
    fn sample_tpe_categorical(
        &self,
        n_choices: usize,
        good_indices: &[usize],
        bad_indices: &[usize],
        rng: &mut StdRng,
    ) -> usize {
        // Count occurrences in good and bad groups
        let mut good_counts = vec![0usize; n_choices];
        let mut bad_counts = vec![0usize; n_choices];

        for &idx in good_indices {
            if idx < n_choices {
                good_counts[idx] += 1;
            }
        }
        for &idx in bad_indices {
            if idx < n_choices {
                bad_counts[idx] += 1;
            }
        }

        // Add smoothing (Laplace smoothing) to avoid zero probabilities
        let good_total = good_indices.len() as f64 + n_choices as f64;
        let bad_total = bad_indices.len() as f64 + n_choices as f64;

        // Calculate l(x)/g(x) ratio for each category
        let mut weights = vec![0.0f64; n_choices];
        for i in 0..n_choices {
            let l_prob = (good_counts[i] as f64 + 1.0) / good_total;
            let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total;
            weights[i] = l_prob / g_prob;
        }

        // Sample proportionally to weights
        let total_weight: f64 = weights.iter().sum();
        let threshold = rng.random::<f64>() * total_weight;

        let mut cumulative = 0.0;
        for (i, &w) in weights.iter().enumerate() {
            cumulative += w;
            if cumulative >= threshold {
                return i;
            }
        }

        // Fallback to last index (shouldn't happen)
        n_choices - 1
    }
}

impl Default for TpeSampler {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for configuring a [`TpeSampler`].
///
/// This builder allows fluent configuration of TPE hyperparameters.
///
/// # Examples
///
/// Using a fixed gamma value:
///
/// ```
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
///
/// let sampler = TpeSamplerBuilder::new()
///     .gamma(0.15)
///     .n_startup_trials(20)
///     .n_ei_candidates(32)
///     .seed(42)
///     .build()
///     .unwrap();
/// ```
///
/// Using a custom gamma strategy:
///
/// ```
/// use optimizer::sampler::tpe::{SqrtGamma, TpeSamplerBuilder};
///
/// let sampler = TpeSamplerBuilder::new()
///     .gamma_strategy(SqrtGamma::default())
///     .n_startup_trials(20)
///     .build()
///     .unwrap();
/// ```
#[derive(Debug, Clone)]
pub struct TpeSamplerBuilder {
    gamma_strategy: Box<dyn GammaStrategy>,
    /// Raw gamma value for deferred validation (Some if `gamma()` was called)
    raw_gamma: Option<f64>,
    n_startup_trials: usize,
    n_ei_candidates: usize,
    kde_bandwidth: Option<f64>,
    seed: Option<u64>,
}

impl TpeSamplerBuilder {
    /// Creates a new builder with default settings.
    ///
    /// Default settings:
    /// - gamma strategy: [`FixedGamma`] with gamma = 0.25
    /// - `n_startup_trials`: 10 (random sampling for first 10 trials)
    /// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample)
    /// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth)
    /// - seed: None (use OS-provided entropy)
    #[must_use]
    pub fn new() -> Self {
        Self {
            gamma_strategy: Box::new(FixedGamma::default()),
            raw_gamma: None,
            n_startup_trials: 10,
            n_ei_candidates: 24,
            kde_bandwidth: None,
            seed: None,
        }
    }

    /// Sets a fixed gamma value for splitting trials into good/bad groups.
    ///
    /// This is a convenience method that creates a [`FixedGamma`] strategy.
    /// For more advanced gamma strategies, use [`gamma_strategy`](Self::gamma_strategy).
    ///
    /// A gamma of 0.25 means the top 25% of trials (by objective value) are
    /// considered "good" and used to build the l(x) distribution.
    ///
    /// # Arguments
    ///
    /// * `gamma` - Quantile value, must be in (0.0, 1.0).
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::TpeSamplerBuilder;
    ///
    /// let sampler = TpeSamplerBuilder::new()
    ///     .gamma(0.10)  // Use top 10% as "good" trials
    ///     .build()
    ///     .unwrap();
    /// ```
    ///
    /// # Note
    ///
    /// Validation happens at `build()` time. If gamma is not in (0.0, 1.0),
    /// `build()` will return `Err(Error::InvalidGamma)`.
    #[must_use]
    pub fn gamma(mut self, gamma: f64) -> Self {
        // We defer validation to build() time for consistency with the existing API
        // Store the raw value for validation later
        self.raw_gamma = Some(gamma);
        self
    }

    /// Sets a custom gamma strategy for splitting trials into good/bad groups.
    ///
    /// The gamma strategy determines what fraction of trials are considered
    /// "good" based on the number of completed trials. This allows the gamma
    /// value to adapt dynamically during optimization.
    ///
    /// # Arguments
    ///
    /// * `strategy` - A type implementing [`GammaStrategy`].
    ///
    /// # Examples
    ///
    /// Using built-in strategies:
    ///
    /// ```
    /// use optimizer::sampler::tpe::{LinearGamma, SqrtGamma, TpeSamplerBuilder};
    ///
    /// // Square root strategy (Optuna-style)
    /// let sampler = TpeSamplerBuilder::new()
    ///     .gamma_strategy(SqrtGamma::default())
    ///     .build()
    ///     .unwrap();
    ///
    /// // Linear interpolation strategy
    /// let sampler = TpeSamplerBuilder::new()
    ///     .gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap())
    ///     .build()
    ///     .unwrap();
    /// ```
    ///
    /// Using a custom strategy:
    ///
    /// ```
    /// use optimizer::sampler::tpe::{GammaStrategy, TpeSamplerBuilder};
    ///
    /// #[derive(Debug, Clone)]
    /// struct MyGamma;
    ///
    /// impl GammaStrategy for MyGamma {
    ///     fn gamma(&self, n_trials: usize) -> f64 {
    ///         0.25 // Always return 0.25
    ///     }
    ///     fn clone_box(&self) -> Box<dyn GammaStrategy> {
    ///         Box::new(self.clone())
    ///     }
    /// }
    ///
    /// let sampler = TpeSamplerBuilder::new()
    ///     .gamma_strategy(MyGamma)
    ///     .build()
    ///     .unwrap();
    /// ```
    #[must_use]
    pub fn gamma_strategy<G: GammaStrategy + 'static>(mut self, strategy: G) -> Self {
        self.gamma_strategy = Box::new(strategy);
        self.raw_gamma = None; // Clear any raw gamma set by gamma()
        self
    }

    /// Sets the number of startup trials before TPE sampling begins.
    ///
    /// During the startup phase, the sampler uses uniform random sampling
    /// to gather initial data. Once `n_startup_trials` have completed,
    /// TPE-based sampling begins.
    ///
    /// # Arguments
    ///
    /// * `n` - Number of random trials before TPE kicks in.
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::TpeSamplerBuilder;
    ///
    /// let sampler = TpeSamplerBuilder::new()
    ///     .n_startup_trials(20)  // Random sample first 20 trials
    ///     .build()
    ///     .unwrap();
    /// ```
    #[must_use]
    pub fn n_startup_trials(mut self, n: usize) -> Self {
        self.n_startup_trials = n;
        self
    }

    /// Sets the number of EI (Expected Improvement) candidates to evaluate.
    ///
    /// When sampling a new point, TPE generates this many candidates from
    /// the l(x) distribution and selects the one with the highest l(x)/g(x)
    /// ratio.
    ///
    /// # Arguments
    ///
    /// * `n` - Number of candidates to evaluate per sample.
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::TpeSamplerBuilder;
    ///
    /// let sampler = TpeSamplerBuilder::new()
    ///     .n_ei_candidates(48)  // Evaluate more candidates
    ///     .build()
    ///     .unwrap();
    /// ```
    #[must_use]
    pub fn n_ei_candidates(mut self, n: usize) -> Self {
        self.n_ei_candidates = n;
        self
    }

    /// Sets a fixed bandwidth for the kernel density estimator.
    ///
    /// By default, TPE uses Scott's rule to automatically select the bandwidth
    /// based on the sample data. Use this method to override with a fixed value.
    ///
    /// Smaller bandwidths give more localized, peaky distributions.
    /// Larger bandwidths give smoother, more spread-out distributions.
    ///
    /// # Arguments
    ///
    /// * `bandwidth` - The fixed bandwidth (standard deviation) for Gaussian kernels.
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::TpeSamplerBuilder;
    ///
    /// let sampler = TpeSamplerBuilder::new()
    ///     .kde_bandwidth(0.5)  // Fixed bandwidth of 0.5
    ///     .build()
    ///     .unwrap();
    /// ```
    ///
    /// # Note
    ///
    /// Validation happens at `build()` time. If bandwidth is not positive,
    /// `build()` will return `Err(Error::InvalidBandwidth)`.
    #[must_use]
    pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self {
        self.kde_bandwidth = Some(bandwidth);
        self
    }

    /// Sets a seed for reproducible sampling.
    ///
    /// # Arguments
    ///
    /// * `seed` - Seed value for the random number generator.
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::TpeSamplerBuilder;
    ///
    /// let sampler = TpeSamplerBuilder::new()
    ///     .seed(42)  // Reproducible results
    ///     .build()
    ///     .unwrap();
    /// ```
    #[must_use]
    pub fn seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }

    /// Builds the configured [`TpeSampler`].
    ///
    /// # Errors
    ///
    /// Returns `Error::InvalidGamma` if a fixed gamma value was set and is not in (0.0, 1.0).
    /// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
    ///
    /// # Examples
    ///
    /// ```
    /// use optimizer::sampler::tpe::TpeSamplerBuilder;
    ///
    /// let sampler = TpeSamplerBuilder::new()
    ///     .gamma(0.15)
    ///     .n_startup_trials(20)
    ///     .n_ei_candidates(32)
    ///     .seed(42)
    ///     .build()
    ///     .unwrap();
    /// ```
    pub fn build(self) -> Result<TpeSampler> {
        // Determine the gamma strategy to use
        let gamma_strategy: Arc<dyn GammaStrategy> = if let Some(raw) = self.raw_gamma {
            // Validate and create FixedGamma from raw value
            Arc::new(FixedGamma::new(raw)?)
        } else {
            Arc::from(self.gamma_strategy)
        };

        // Validate bandwidth
        if let Some(bw) = self.kde_bandwidth
            && bw <= 0.0
        {
            return Err(Error::InvalidBandwidth(bw));
        }

        let rng = match self.seed {
            Some(s) => StdRng::seed_from_u64(s),
            None => StdRng::from_os_rng(),
        };

        Ok(TpeSampler {
            gamma_strategy,
            n_startup_trials: self.n_startup_trials,
            n_ei_candidates: self.n_ei_candidates,
            kde_bandwidth: self.kde_bandwidth,
            rng: Mutex::new(rng),
        })
    }
}

impl Default for TpeSamplerBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl Sampler for TpeSampler {
    #[allow(clippy::too_many_lines)]
    fn sample(
        &self,
        distribution: &Distribution,
        _trial_id: u64,
        history: &[CompletedTrial],
    ) -> ParamValue {
        let mut rng = self.rng.lock();

        // Fall back to random sampling during startup phase
        if history.len() < self.n_startup_trials {
            return self.sample_uniform(distribution, &mut rng);
        }

        // Split trials into good and bad groups
        let (good_trials, bad_trials) = self.split_trials(history);

        // Need at least 1 trial in each group for TPE
        if good_trials.is_empty() || bad_trials.is_empty() {
            return self.sample_uniform(distribution, &mut rng);
        }

        // Extract parameter values for this distribution
        // Since we don't have the parameter name here, we need to look at all
        // trials and find matching distributions
        // Note: This is a simplification - in practice, we'd need the param name
        // For now, we'll collect values from trials that have this exact distribution type

        match distribution {
            Distribution::Float(d) => {
                // Collect float values from trials
                let good_values: Vec<f64> = good_trials
                    .iter()
                    .flat_map(|t| t.params.values())
                    .filter_map(|v| match v {
                        ParamValue::Float(f) => Some(*f),
                        _ => None,
                    })
                    .filter(|&v| v >= d.low && v <= d.high)
                    .collect();

                let bad_values: Vec<f64> = bad_trials
                    .iter()
                    .flat_map(|t| t.params.values())
                    .filter_map(|v| match v {
                        ParamValue::Float(f) => Some(*f),
                        _ => None,
                    })
                    .filter(|&v| v >= d.low && v <= d.high)
                    .collect();

                // Need values in both groups for TPE
                if good_values.is_empty() || bad_values.is_empty() {
                    return self.sample_uniform(distribution, &mut rng);
                }

                let value = self.sample_tpe_float(
                    d.low,
                    d.high,
                    d.log_scale,
                    d.step,
                    good_values,
                    bad_values,
                    &mut rng,
                );
                ParamValue::Float(value)
            }
            Distribution::Int(d) => {
                let good_values: Vec<i64> = good_trials
                    .iter()
                    .flat_map(|t| t.params.values())
                    .filter_map(|v| match v {
                        ParamValue::Int(i) => Some(*i),
                        _ => None,
                    })
                    .filter(|&v| v >= d.low && v <= d.high)
                    .collect();

                let bad_values: Vec<i64> = bad_trials
                    .iter()
                    .flat_map(|t| t.params.values())
                    .filter_map(|v| match v {
                        ParamValue::Int(i) => Some(*i),
                        _ => None,
                    })
                    .filter(|&v| v >= d.low && v <= d.high)
                    .collect();

                if good_values.is_empty() || bad_values.is_empty() {
                    return self.sample_uniform(distribution, &mut rng);
                }

                let value = self.sample_tpe_int(
                    d.low,
                    d.high,
                    d.log_scale,
                    d.step,
                    &good_values,
                    &bad_values,
                    &mut rng,
                );
                ParamValue::Int(value)
            }
            Distribution::Categorical(d) => {
                let good_indices: Vec<usize> = good_trials
                    .iter()
                    .flat_map(|t| t.params.values())
                    .filter_map(|v| match v {
                        ParamValue::Categorical(i) => Some(*i),
                        _ => None,
                    })
                    .filter(|&i| i < d.n_choices)
                    .collect();

                let bad_indices: Vec<usize> = bad_trials
                    .iter()
                    .flat_map(|t| t.params.values())
                    .filter_map(|v| match v {
                        ParamValue::Categorical(i) => Some(*i),
                        _ => None,
                    })
                    .filter(|&i| i < d.n_choices)
                    .collect();

                if good_indices.is_empty() || bad_indices.is_empty() {
                    return self.sample_uniform(distribution, &mut rng);
                }

                let index =
                    self.sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, &mut rng);
                ParamValue::Categorical(index)
            }
        }
    }
}

#[cfg(test)]
#[allow(
    clippy::similar_names,
    clippy::cast_sign_loss,
    clippy::cast_precision_loss
)]
mod tests {
    use std::collections::HashMap;

    use super::*;
    use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};

    fn create_trial(
        id: u64,
        value: f64,
        params: Vec<(&str, ParamValue, Distribution)>,
    ) -> CompletedTrial {
        let mut param_map = HashMap::new();
        let mut dist_map = HashMap::new();
        for (name, pv, dist) in params {
            param_map.insert(name.to_string(), pv);
            dist_map.insert(name.to_string(), dist);
        }
        CompletedTrial::new(id, param_map, dist_map, value)
    }

    #[test]
    fn test_tpe_sampler_new() {
        let sampler = TpeSampler::new();
        // Default uses FixedGamma with 0.25
        assert!((sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON);
        assert_eq!(sampler.n_startup_trials, 10);
        assert_eq!(sampler.n_ei_candidates, 24);
    }

    #[test]
    fn test_tpe_sampler_with_config() {
        let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42)).unwrap();
        // with_config uses FixedGamma
        assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
        assert_eq!(sampler.n_startup_trials, 20);
        assert_eq!(sampler.n_ei_candidates, 32);
    }

    #[test]
    fn test_tpe_sampler_invalid_gamma_zero() {
        let result = TpeSampler::with_config(0.0, 10, 24, None, None);
        assert!(matches!(result, Err(Error::InvalidGamma(_))));
    }

    #[test]
    fn test_tpe_sampler_invalid_gamma_one() {
        let result = TpeSampler::with_config(1.0, 10, 24, None, None);
        assert!(matches!(result, Err(Error::InvalidGamma(_))));
    }

    #[test]
    fn test_tpe_startup_random_sampling() {
        let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)).unwrap();
        let dist = Distribution::Float(FloatDistribution {
            low: 0.0,
            high: 1.0,
            log_scale: false,
            step: None,
        });

        // With fewer than n_startup_trials, should use random sampling
        let history: Vec<CompletedTrial> = vec![];

        for _ in 0..100 {
            let value = sampler.sample(&dist, 0, &history);
            if let ParamValue::Float(v) = value {
                assert!((0.0..=1.0).contains(&v));
            } else {
                panic!("Expected Float value");
            }
        }
    }

    #[test]
    fn test_tpe_split_trials() {
        let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)).unwrap();

        let dist = Distribution::Float(FloatDistribution {
            low: 0.0,
            high: 1.0,
            log_scale: false,
            step: None,
        });

        // Create 20 trials with values 0..20
        let history: Vec<CompletedTrial> = (0..20)
            .map(|i| {
                create_trial(
                    i as u64,
                    f64::from(i),
                    vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
                )
            })
            .collect();

        let (good, bad) = sampler.split_trials(&history);

        // With gamma=0.25 and 20 trials, should have 5 good and 15 bad
        assert_eq!(good.len(), 5);
        assert_eq!(bad.len(), 15);

        // Good trials should have lowest values
        for trial in &good {
            assert!(trial.value < 5.0);
        }
    }

    #[test]
    fn test_tpe_samples_float_with_history() {
        let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();

        let dist = Distribution::Float(FloatDistribution {
            low: 0.0,
            high: 1.0,
            log_scale: false,
            step: None,
        });

        // Create history where low values (near 0.2) are "good"
        let history: Vec<CompletedTrial> = (0..20)
            .map(|i| {
                let x = f64::from(i) / 20.0;
                // Objective is (x - 0.2)^2, minimized at x=0.2
                let value = (x - 0.2).powi(2);
                create_trial(
                    i as u64,
                    value,
                    vec![("x", ParamValue::Float(x), dist.clone())],
                )
            })
            .collect();

        // TPE should bias toward values near 0.2
        let mut samples = vec![];
        for i in 0..100 {
            let value = sampler.sample(&dist, 100 + i, &history);
            if let ParamValue::Float(v) = value {
                samples.push(v);
            }
        }

        // Calculate mean of samples - should be closer to 0.2 than 0.5
        let mean: f64 = samples.iter().sum::<f64>() / samples.len() as f64;
        assert!(
            mean < 0.5,
            "Mean {mean} should be less than 0.5 (biased toward good region near 0.2)"
        );
    }

    #[test]
    fn test_tpe_categorical_sampling() {
        let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();

        let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 });

        // Create history where category 1 is consistently good
        let history: Vec<CompletedTrial> = (0..20)
            .map(|i| {
                let category = i % 4;
                // Category 1 has best (lowest) objective value
                let value = if category == 1 { 0.0 } else { 1.0 };
                create_trial(
                    i as u64,
                    value,
                    vec![(
                        "cat",
                        ParamValue::Categorical(category as usize),
                        dist.clone(),
                    )],
                )
            })
            .collect();

        // TPE should favor category 1
        let mut counts = vec![0usize; 4];
        for i in 0..100 {
            let value = sampler.sample(&dist, 100 + i, &history);
            if let ParamValue::Categorical(idx) = value {
                counts[idx] += 1;
            }
        }

        // Category 1 should be sampled more often
        assert!(
            counts[1] > counts[0] && counts[1] > counts[2] && counts[1] > counts[3],
            "Category 1 should be most common: {counts:?}"
        );
    }

    #[test]
    fn test_tpe_int_sampling() {
        let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();

        let dist = Distribution::Int(IntDistribution {
            low: 0,
            high: 100,
            log_scale: false,
            step: None,
        });

        // Create history where values near 30 are good
        let history: Vec<CompletedTrial> = (0..20)
            .map(|i| {
                let x = i * 5; // 0, 5, 10, ..., 95
                let value = ((x as f64) - 30.0).powi(2);
                create_trial(
                    i as u64,
                    value,
                    vec![("x", ParamValue::Int(x), dist.clone())],
                )
            })
            .collect();

        // TPE should bias toward values near 30
        for i in 0..50 {
            let value = sampler.sample(&dist, 100 + i, &history);
            if let ParamValue::Int(v) = value {
                assert!((0..=100).contains(&v), "Value {v} out of range");
            } else {
                panic!("Expected Int value");
            }
        }
    }

    #[test]
    fn test_tpe_reproducibility() {
        let dist = Distribution::Float(FloatDistribution {
            low: 0.0,
            high: 1.0,
            log_scale: false,
            step: None,
        });

        let history: Vec<CompletedTrial> = (0..20)
            .map(|i| {
                create_trial(
                    i as u64,
                    f64::from(i),
                    vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
                )
            })
            .collect();

        let sampler1 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345)).unwrap();
        let sampler2 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345)).unwrap();

        for i in 0..10 {
            let v1 = sampler1.sample(&dist, i, &history);
            let v2 = sampler2.sample(&dist, i, &history);
            assert_eq!(v1, v2, "Samples should be identical with same seed");
        }
    }

    #[test]
    fn test_tpe_sampler_builder_default() {
        let builder = TpeSamplerBuilder::new();
        let sampler = builder.build().unwrap();
        assert!((sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON);
        assert_eq!(sampler.n_startup_trials, 10);
        assert_eq!(sampler.n_ei_candidates, 24);
    }

    #[test]
    fn test_tpe_sampler_builder_custom() {
        let sampler = TpeSamplerBuilder::new()
            .gamma(0.15)
            .n_startup_trials(20)
            .n_ei_candidates(32)
            .seed(42)
            .build()
            .unwrap();
        assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
        assert_eq!(sampler.n_startup_trials, 20);
        assert_eq!(sampler.n_ei_candidates, 32);
    }

    #[test]
    fn test_tpe_sampler_builder_via_sampler() {
        let sampler = TpeSampler::builder()
            .gamma(0.10)
            .n_startup_trials(15)
            .n_ei_candidates(48)
            .build()
            .unwrap();
        assert!((sampler.gamma_strategy().gamma(0) - 0.10).abs() < f64::EPSILON);
        assert_eq!(sampler.n_startup_trials, 15);
        assert_eq!(sampler.n_ei_candidates, 48);
    }

    #[test]
    fn test_tpe_sampler_builder_partial() {
        // Test setting only some options
        let sampler = TpeSamplerBuilder::new().gamma(0.20).build().unwrap();
        assert!((sampler.gamma_strategy().gamma(0) - 0.20).abs() < f64::EPSILON);
        assert_eq!(sampler.n_startup_trials, 10); // default
        assert_eq!(sampler.n_ei_candidates, 24); // default
    }

    #[test]
    fn test_tpe_sampler_builder_invalid_gamma() {
        let result = TpeSamplerBuilder::new().gamma(1.5).build();
        assert!(matches!(result, Err(Error::InvalidGamma(_))));
    }

    #[test]
    fn test_tpe_sampler_builder_reproducibility() {
        let dist = Distribution::Float(FloatDistribution {
            low: 0.0,
            high: 1.0,
            log_scale: false,
            step: None,
        });

        let history: Vec<CompletedTrial> = (0..20u32)
            .map(|i| {
                create_trial(
                    u64::from(i),
                    f64::from(i),
                    vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
                )
            })
            .collect();

        let sampler1 = TpeSampler::builder()
            .seed(99999)
            .n_startup_trials(5)
            .build()
            .unwrap();
        let sampler2 = TpeSampler::builder()
            .seed(99999)
            .n_startup_trials(5)
            .build()
            .unwrap();

        for i in 0..10 {
            let v1 = sampler1.sample(&dist, i, &history);
            let v2 = sampler2.sample(&dist, i, &history);
            assert_eq!(
                v1, v2,
                "Builder-created samplers with same seed should be identical"
            );
        }
    }

    // ========================================================================
    // Gamma Strategy Tests
    // ========================================================================

    #[test]
    fn test_fixed_gamma_default() {
        let strategy = FixedGamma::default();
        assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);
        assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);
        assert!((strategy.value() - 0.25).abs() < f64::EPSILON);
    }

    #[test]
    fn test_fixed_gamma_custom() {
        let strategy = FixedGamma::new(0.15).unwrap();
        assert!((strategy.gamma(0) - 0.15).abs() < f64::EPSILON);
        assert!((strategy.gamma(50) - 0.15).abs() < f64::EPSILON);
        assert!((strategy.gamma(1000) - 0.15).abs() < f64::EPSILON);
    }

    #[test]
    fn test_fixed_gamma_invalid() {
        assert!(FixedGamma::new(0.0).is_err());
        assert!(FixedGamma::new(1.0).is_err());
        assert!(FixedGamma::new(-0.1).is_err());
        assert!(FixedGamma::new(1.5).is_err());
    }

    #[test]
    fn test_linear_gamma_default() {
        let strategy = LinearGamma::default();
        assert!((strategy.gamma(0) - 0.10).abs() < f64::EPSILON);
        assert!((strategy.gamma(50) - 0.175).abs() < f64::EPSILON); // midpoint
        assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);
        assert!((strategy.gamma(200) - 0.25).abs() < f64::EPSILON); // capped
    }

    #[test]
    fn test_linear_gamma_custom() {
        let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap();
        assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON);
        assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON);
        assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON);
        assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON);
    }

    #[test]
    fn test_linear_gamma_invalid() {
        assert!(LinearGamma::new(0.0, 0.5, 100).is_err());
        assert!(LinearGamma::new(0.1, 1.0, 100).is_err());
        assert!(LinearGamma::new(0.5, 0.2, 100).is_err()); // min > max
    }

    #[test]
    fn test_sqrt_gamma_default() {
        let strategy = SqrtGamma::default();
        // At n=0, returns gamma_max
        assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);

        // gamma decreases with more trials
        let g10 = strategy.gamma(10);
        let g100 = strategy.gamma(100);
        assert!(g10 > g100);
    }

    #[test]
    fn test_sqrt_gamma_custom() {
        let strategy = SqrtGamma::new(2.0, 0.5).unwrap();
        assert!((strategy.gamma(0) - 0.5).abs() < f64::EPSILON);

        // At n=4: n_good = max(1, 2/2) = 1, gamma = 1/4 = 0.25
        let g4 = strategy.gamma(4);
        assert!((g4 - 0.25).abs() < f64::EPSILON);
    }

    #[test]
    fn test_sqrt_gamma_invalid() {
        assert!(SqrtGamma::new(0.0, 0.25).is_err()); // factor must be positive
        assert!(SqrtGamma::new(-1.0, 0.25).is_err());
        assert!(SqrtGamma::new(1.0, 0.0).is_err());
        assert!(SqrtGamma::new(1.0, 1.0).is_err());
    }

    #[test]
    fn test_hyperopt_gamma_default() {
        let strategy = HyperoptGamma::default();
        // At n=0, returns gamma_max
        assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);

        // At n=100: (24+1)/100 = 0.25, so capped to 0.25
        assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);

        // At n=200: (24+1)/200 = 0.125
        assert!((strategy.gamma(200) - 0.125).abs() < f64::EPSILON);
    }

    #[test]
    fn test_hyperopt_gamma_custom() {
        let strategy = HyperoptGamma::new(9.0, 0.5).unwrap();
        // At n=20: (9+1)/20 = 0.5, capped to 0.5
        assert!((strategy.gamma(20) - 0.5).abs() < f64::EPSILON);

        // At n=100: (9+1)/100 = 0.1
        assert!((strategy.gamma(100) - 0.1).abs() < f64::EPSILON);
    }

    #[test]
    fn test_hyperopt_gamma_invalid() {
        assert!(HyperoptGamma::new(-1.0, 0.25).is_err());
        assert!(HyperoptGamma::new(24.0, 0.0).is_err());
        assert!(HyperoptGamma::new(24.0, 1.0).is_err());
    }

    #[test]
    fn test_gamma_strategy_clone_box() {
        let fixed: Box<dyn GammaStrategy> = Box::new(FixedGamma::new(0.3).unwrap());
        let cloned = fixed.clone();
        assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON);

        let linear: Box<dyn GammaStrategy> = Box::new(LinearGamma::default());
        let cloned = linear.clone();
        assert!((cloned.gamma(0) - 0.10).abs() < f64::EPSILON);
    }

    #[test]
    fn test_tpe_with_sqrt_gamma_strategy() {
        let sampler = TpeSampler::builder()
            .gamma_strategy(SqrtGamma::default())
            .n_startup_trials(5)
            .seed(42)
            .build()
            .unwrap();

        let dist = Distribution::Float(FloatDistribution {
            low: 0.0,
            high: 1.0,
            log_scale: false,
            step: None,
        });

        let history: Vec<CompletedTrial> = (0..20)
            .map(|i| {
                create_trial(
                    i as u64,
                    f64::from(i),
                    vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
                )
            })
            .collect();

        // Should be able to sample with the sqrt gamma strategy
        let value = sampler.sample(&dist, 100, &history);
        if let ParamValue::Float(v) = value {
            assert!((0.0..=1.0).contains(&v));
        } else {
            panic!("Expected Float value");
        }
    }

    #[test]
    fn test_tpe_with_linear_gamma_strategy() {
        let sampler = TpeSampler::builder()
            .gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap())
            .n_startup_trials(5)
            .seed(42)
            .build()
            .unwrap();

        // Verify the strategy is applied
        let g = sampler.gamma_strategy().gamma(25);
        assert!((g - 0.2).abs() < f64::EPSILON); // midpoint of 0.1 to 0.3
    }

    #[test]
    fn test_tpe_with_hyperopt_gamma_strategy() {
        let sampler = TpeSampler::builder()
            .gamma_strategy(HyperoptGamma::default())
            .n_startup_trials(5)
            .seed(42)
            .build()
            .unwrap();

        let dist = Distribution::Float(FloatDistribution {
            low: 0.0,
            high: 1.0,
            log_scale: false,
            step: None,
        });

        let history: Vec<CompletedTrial> = (0..20)
            .map(|i| {
                create_trial(
                    i as u64,
                    f64::from(i),
                    vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
                )
            })
            .collect();

        // Should be able to sample with the hyperopt gamma strategy
        let value = sampler.sample(&dist, 100, &history);
        if let ParamValue::Float(v) = value {
            assert!((0.0..=1.0).contains(&v));
        } else {
            panic!("Expected Float value");
        }
    }

    #[test]
    fn test_gamma_overrides_gamma_strategy() {
        // When gamma() is called after gamma_strategy(), it should take precedence
        let sampler = TpeSampler::builder()
            .gamma_strategy(SqrtGamma::default())
            .gamma(0.15) // This should override
            .build()
            .unwrap();

        // Should use fixed gamma of 0.15
        assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
        assert!((sampler.gamma_strategy().gamma(100) - 0.15).abs() < f64::EPSILON);
    }

    #[test]
    fn test_gamma_strategy_overrides_gamma() {
        // When gamma_strategy() is called after gamma(), it should take precedence
        let sampler = TpeSampler::builder()
            .gamma(0.15)
            .gamma_strategy(SqrtGamma::default()) // This should override
            .build()
            .unwrap();

        // Should use SqrtGamma - gamma decreases with trials
        let g10 = sampler.gamma_strategy().gamma(10);
        let g100 = sampler.gamma_strategy().gamma(100);
        assert!(g10 > g100, "SqrtGamma should decrease with more trials");
    }

    #[test]
    fn test_with_strategy_constructor() {
        let sampler = TpeSampler::with_strategy(
            LinearGamma::new(0.1, 0.4, 100).unwrap(),
            15,
            32,
            None,
            Some(42),
        )
        .unwrap();

        assert_eq!(sampler.n_startup_trials, 15);
        assert_eq!(sampler.n_ei_candidates, 32);
        assert!((sampler.gamma_strategy().gamma(0) - 0.1).abs() < f64::EPSILON);
        assert!((sampler.gamma_strategy().gamma(100) - 0.4).abs() < f64::EPSILON);
    }

    #[test]
    fn test_custom_gamma_strategy() {
        #[derive(Debug, Clone)]
        struct DoubleGamma;

        impl GammaStrategy for DoubleGamma {
            fn gamma(&self, n_trials: usize) -> f64 {
                // Double the trial count-based calculation, capped at 0.5
                (0.01 * n_trials as f64).min(0.5)
            }

            fn clone_box(&self) -> Box<dyn GammaStrategy> {
                Box::new(self.clone())
            }
        }

        let sampler = TpeSampler::builder()
            .gamma_strategy(DoubleGamma)
            .build()
            .unwrap();

        assert!((sampler.gamma_strategy().gamma(10) - 0.1).abs() < f64::EPSILON);
        assert!((sampler.gamma_strategy().gamma(50) - 0.5).abs() < f64::EPSILON);
        assert!((sampler.gamma_strategy().gamma(100) - 0.5).abs() < f64::EPSILON);
    }
}