//! Tree-Parzen Estimator (TPE) for Bayesian Hyperparameter Optimization
//!
//! TPE builds separate density models for "good" and "bad" trial parameters,
//! then suggests new points by maximizing Expected Improvement = l(x)/g(x).
//!
//! ## Algorithm
//!
//! 1. Run N initial random trials (Latin Hypercube Sampling)
//! 2. After enough data, split trials into "good" (top gamma=25%) and "bad" (bottom 75%)
//! 3. Build separate per-dimension kernel density estimates (KDEs) for good and bad groups
//! 4. Suggest next point by sampling from the "good" KDE and scoring by EI = l(x)/g(x)
//! 5. The optimizer MINIMIZES the objective (lower = better)
//!
//! ## References
//!
//! Bergstra, J., Bardenet, R., Bengio, Y., & Kegl, B. (2011).
//! "Algorithms for Hyper-Parameter Optimization." NeurIPS.

use rand::prelude::*;
use rand::SeedableRng;
use serde::{Deserialize, Serialize};
use std::path::Path;

/// A single completed trial recording parameters and objective value.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrialRecord {
    /// Parameter values in continuous space.
    pub params: Vec<f64>,
    /// Objective value (lower is better).
    pub objective: f64,
}

/// TPE optimizer configuration.
#[derive(Debug, Clone)]
pub struct TpeConfig {
    /// Number of dimensions in parameter space.
    pub n_dims: usize,
    /// Maximum number of trials.
    pub max_trials: usize,
    /// Number of initial LHS samples before TPE kicks in.
    pub n_initial: usize,
    /// Quantile for good/bad split (0.25 = top 25% are "good").
    pub gamma: f64,
    /// Number of candidates to evaluate for EI.
    pub n_candidates: usize,
    /// Random seed (optional).
    pub seed: Option<u64>,
}

impl TpeConfig {
    /// Create a new TPE configuration with sensible defaults.
    ///
    /// - `n_initial` is max(5, n_dims) to ensure enough initial exploration.
    /// - `gamma` defaults to 0.25 (top 25% are "good").
    /// - `n_candidates` defaults to 100 EI candidates per suggestion.
    pub fn new(n_dims: usize, max_trials: usize) -> Self {
        Self {
            n_dims,
            max_trials,
            n_initial: 5.max(n_dims),
            gamma: 0.25,
            n_candidates: 100,
            seed: None,
        }
    }
}

/// Tree-Parzen Estimator optimizer.
///
/// Maintains a history of completed trials and uses kernel density estimation
/// to suggest promising new parameter configurations.
#[derive(Debug)]
pub struct TpeOptimizer {
    config: TpeConfig,
    /// Completed trial history.
    pub trials: Vec<TrialRecord>,
    rng: StdRng,
}

impl TpeOptimizer {
    /// Create a new TPE optimizer from configuration.
    pub fn new(config: TpeConfig) -> Self {
        let rng = match config.seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => StdRng::from_entropy(),
        };
        Self {
            config,
            trials: Vec::new(),
            rng,
        }
    }

    /// Suggest the next parameter configuration to evaluate.
    ///
    /// During the initial phase (fewer trials than `n_initial`), returns a
    /// Latin Hypercube sample. After that, uses TPE to suggest the point
    /// with the highest expected improvement.
    pub fn suggest(&mut self, bounds: &[(f64, f64)]) -> Vec<f64> {
        if self.trials.len() < self.config.n_initial {
            // Initial exploration phase: return single LHS sample
            let mut samples = self.latin_hypercube_sample(bounds, 1);
            return samples.pop().unwrap_or_else(|| {
                // Fallback: uniform random sample
                bounds.iter().map(|&(lo, hi)| self.rng.gen_range(lo..=hi)).collect()
            });
        }

        // TPE phase: split, build KDEs, sample candidates, pick best EI
        // Clone parameter vectors to break the borrow on self.trials so we can
        // mutably borrow self.rng via sample_from_kde.
        let (good, bad) = self.split_trials();
        let good_params: Vec<Vec<f64>> = good.iter().map(|t| t.params.clone()).collect();
        let bad_params: Vec<Vec<f64>> = bad.iter().map(|t| t.params.clone()).collect();
        let good_refs: Vec<&Vec<f64>> = good_params.iter().collect();
        let bad_refs: Vec<&Vec<f64>> = bad_params.iter().collect();

        // Generate candidates by sampling from the good KDE
        let mut best_candidate = self.sample_from_kde(&good_refs, bounds);
        let mut best_ei = Self::expected_improvement_static(&best_candidate, &good_refs, &bad_refs, bounds);

        for _ in 1..self.config.n_candidates {
            let candidate = self.sample_from_kde(&good_refs, bounds);
            let ei = Self::expected_improvement_static(&candidate, &good_refs, &bad_refs, bounds);
            if ei > best_ei {
                best_ei = ei;
                best_candidate = candidate;
            }
        }

        // Best-trial injection: always evaluate EI at (and around) the best known point.
        // Prevents the optimizer from "forgetting" a good trial when KDE is too wide.
        let best_params_opt: Option<Vec<f64>> = self
            .trials
            .iter()
            .min_by(|a, b| {
                a.objective
                    .partial_cmp(&b.objective)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|t| t.params.clone());

        if let Some(best_params) = best_params_opt {
            let ei = Self::expected_improvement_static(&best_params, &good_refs, &bad_refs, bounds);
            if ei > best_ei {
                best_ei = ei;
                best_candidate = best_params.clone();
            }
            // Small perturbations (±5% per dimension) around best trial
            for _ in 0..5 {
                let mut perturbed = best_params.clone();
                for (d, val) in perturbed.iter_mut().enumerate() {
                    if let Some(&(lo, hi)) = bounds.get(d) {
                        let noise = self.rng.gen_range(-0.05..0.05) * (hi - lo);
                        *val = (*val + noise).clamp(lo, hi);
                    }
                }
                let ei =
                    Self::expected_improvement_static(&perturbed, &good_refs, &bad_refs, bounds);
                if ei > best_ei {
                    best_ei = ei;
                    best_candidate = perturbed;
                }
            }
        }

        best_candidate
    }

    /// Record a completed trial.
    pub fn add_trial(&mut self, params: Vec<f64>, objective: f64) {
        self.trials.push(TrialRecord { params, objective });
    }

    /// Split trials into "good" (top gamma%) and "bad" (bottom 1-gamma%).
    ///
    /// Trials are sorted by objective ascending (lower = better), so the first
    /// `ceil(n * gamma)` trials are "good".
    pub fn split_trials(&self) -> (Vec<&TrialRecord>, Vec<&TrialRecord>) {
        let mut sorted: Vec<&TrialRecord> = self.trials.iter().collect();
        sorted.sort_by(|a, b| {
            a.objective
                .partial_cmp(&b.objective)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        let n_good = (sorted.len() as f64 * self.config.gamma).ceil() as usize;
        let n_good = n_good.max(1).min(sorted.len().saturating_sub(1));

        let (good, bad) = sorted.split_at(n_good);
        (good.to_vec(), bad.to_vec())
    }

    /// Generate `n` Latin Hypercube samples within the given bounds.
    ///
    /// For each dimension, divides [0, 1] into `n` equal bins, shuffles them,
    /// and scales to the parameter bounds.
    pub fn latin_hypercube_sample(&mut self, bounds: &[(f64, f64)], n: usize) -> Vec<Vec<f64>> {
        if n == 0 || bounds.is_empty() {
            return Vec::new();
        }

        let n_dims = bounds.len();
        let mut samples = vec![vec![0.0; n_dims]; n];

        for (d, &(lo, hi)) in bounds.iter().enumerate() {
            // Create permutation of bin indices
            let mut indices: Vec<usize> = (0..n).collect();
            indices.shuffle(&mut self.rng);

            for (i, &idx) in indices.iter().enumerate() {
                // Sample uniformly within the bin
                let bin_lo = idx as f64 / n as f64;
                let bin_hi = (idx + 1) as f64 / n as f64;
                let u = self.rng.gen_range(bin_lo..bin_hi);
                // Scale to bounds
                if let Some(s) = samples.get_mut(i) {
                    if let Some(v) = s.get_mut(d) {
                        *v = lo + u * (hi - lo);
                    }
                }
            }
        }

        samples
    }

    /// Compute the log-PDF of a point under an independent per-dimension KDE.
    ///
    /// Uses Gaussian kernels with Silverman's rule for bandwidth selection.
    /// Returns the sum of log-PDFs across dimensions (independence assumption).
    /// Uses log-sum-exp for numerical stability.
    pub fn kde_log_pdf(samples: &[&Vec<f64>], point: &[f64], bounds: &[(f64, f64)]) -> f64 {
        if samples.is_empty() {
            return f64::NEG_INFINITY;
        }

        let n = samples.len() as f64;
        let n_dims = point.len();
        let mut total_log_pdf = 0.0;

        for d in 0..n_dims {
            // Collect values for this dimension
            let values: Vec<f64> = samples
                .iter()
                .filter_map(|s| s.get(d).copied())
                .collect();

            if values.is_empty() {
                continue;
            }

            // Scott's rule bandwidth (dimension-aware, tighter in high-D)
            let bandwidth = Self::scott_bandwidth(&values, bounds.get(d).map(|&(lo, hi)| hi - lo), n_dims);

            // Compute log-PDF at point[d] using log-sum-exp
            let Some(&x) = point.get(d) else {
                continue;
            };

            // log-sum-exp: log(sum(exp(log_k_i))) where log_k_i = -0.5 * ((x - mu_i) / h)^2 - log(h) - 0.5*log(2pi)
            let log_norm = -0.5_f64 * (2.0 * std::f64::consts::PI).ln() - bandwidth.ln();
            let log_kernels: Vec<f64> = values
                .iter()
                .map(|&mu| {
                    let z = (x - mu) / bandwidth;
                    log_norm - 0.5 * z * z
                })
                .collect();

            let log_pdf_d = Self::log_sum_exp(&log_kernels) - n.ln();
            total_log_pdf += log_pdf_d;
        }

        total_log_pdf
    }

    /// Compute the expected improvement score (in log space) for a candidate point.
    ///
    /// EI = l(x) / g(x), so log(EI) = log_l(x) - log_g(x).
    /// Higher is better.
    pub fn expected_improvement(
        &self,
        point: &[f64],
        good_params: &[&Vec<f64>],
        bad_params: &[&Vec<f64>],
        bounds: &[(f64, f64)],
    ) -> f64 {
        Self::expected_improvement_static(point, good_params, bad_params, bounds)
    }

    /// Static version of expected_improvement (no &self borrow needed).
    fn expected_improvement_static(
        point: &[f64],
        good_params: &[&Vec<f64>],
        bad_params: &[&Vec<f64>],
        bounds: &[(f64, f64)],
    ) -> f64 {
        let log_l = Self::kde_log_pdf(good_params, point, bounds);
        let log_g = Self::kde_log_pdf(bad_params, point, bounds);
        log_l - log_g
    }

    /// Save trial history to a JSON file.
    pub fn save_history(&self, path: &Path) -> Result<(), std::io::Error> {
        let json = serde_json::to_string_pretty(&self.trials).map_err(|e| {
            std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())
        })?;
        std::fs::write(path, json)
    }

    /// Load trial history from a JSON file.
    ///
    /// Returns the number of trials loaded. Loaded trials are appended to
    /// any existing trials in the optimizer.
    pub fn load_history(&mut self, path: &Path) -> Result<usize, std::io::Error> {
        let json = std::fs::read_to_string(path)?;
        let records: Vec<TrialRecord> = serde_json::from_str(&json).map_err(|e| {
            std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())
        })?;
        let count = records.len();
        self.trials.extend(records);
        Ok(count)
    }

    /// Return a reference to the best trial (lowest objective), if any.
    pub fn best_trial(&self) -> Option<&TrialRecord> {
        self.trials
            .iter()
            .min_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap_or(std::cmp::Ordering::Equal))
    }

    /// Return the number of completed trials.
    pub fn n_trials(&self) -> usize {
        self.trials.len()
    }

    // --- Private helpers ---

    /// Compute KDE bandwidth using Scott's rule (dimension-aware).
    ///
    /// h = 0.7 * σ * n^(-1/(d+4)), tighter than Silverman's in high-D spaces.
    /// Falls back to range/100 floor when all samples are identical.
    fn scott_bandwidth(values: &[f64], range: Option<f64>, n_dims: usize) -> f64 {
        let n = values.len() as f64;
        if n < 1.0 {
            return 1.0;
        }

        let mean = values.iter().sum::<f64>() / n;
        let variance = values.iter().map(|&v| (v - mean) * (v - mean)).sum::<f64>() / n;
        let std_dev = variance.sqrt();

        // Scott's rule with 0.7 prefactor for tighter exploitation
        let exponent = -1.0 / (n_dims as f64 + 4.0);
        let bandwidth = 0.7 * std_dev * n.powf(exponent);

        // Floor: avoid zero bandwidth (when all samples are the same)
        let floor = match range {
            Some(r) if r > 0.0 => r / 100.0,
            _ => 0.01,
        };

        bandwidth.max(floor)
    }

    /// Numerically stable log-sum-exp: log(sum(exp(values))).
    fn log_sum_exp(values: &[f64]) -> f64 {
        if values.is_empty() {
            return f64::NEG_INFINITY;
        }

        let max_val = values
            .iter()
            .copied()
            .fold(f64::NEG_INFINITY, f64::max);

        if max_val == f64::NEG_INFINITY {
            return f64::NEG_INFINITY;
        }

        let sum_exp: f64 = values.iter().map(|&v| (v - max_val).exp()).sum();
        max_val + sum_exp.ln()
    }

    /// Sample a point from a KDE defined by the given samples.
    ///
    /// For each dimension independently:
    /// 1. Pick a random sample (kernel center)
    /// 2. Add Gaussian noise with Scott bandwidth
    /// 3. Clamp to bounds
    fn sample_from_kde(&mut self, samples: &[&Vec<f64>], bounds: &[(f64, f64)]) -> Vec<f64> {
        let n_dims = bounds.len();
        let mut point = vec![0.0; n_dims];

        if samples.is_empty() {
            // Fallback: uniform random
            for (d, &(lo, hi)) in bounds.iter().enumerate() {
                if let Some(v) = point.get_mut(d) {
                    *v = self.rng.gen_range(lo..=hi);
                }
            }
            return point;
        }

        for (d, &(lo, hi)) in bounds.iter().enumerate() {
            // Collect values for this dimension
            let values: Vec<f64> = samples
                .iter()
                .filter_map(|s| s.get(d).copied())
                .collect();

            if values.is_empty() {
                if let Some(v) = point.get_mut(d) {
                    *v = self.rng.gen_range(lo..=hi);
                }
                continue;
            }

            let bandwidth = Self::scott_bandwidth(&values, Some(hi - lo), n_dims);

            // Pick a random kernel center
            let center_idx = self.rng.gen_range(0..values.len());
            let center = values.get(center_idx).copied().unwrap_or((lo + hi) / 2.0);

            // Sample from Gaussian kernel and clamp to bounds
            let noise: f64 = self.rng.sample::<f64, _>(rand::distributions::Standard) * bandwidth;
            let raw = center + noise;
            if let Some(v) = point.get_mut(d) {
                *v = raw.clamp(lo, hi);
            }
        }

        point
    }
}

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

    #[test]
    fn test_tpe_split_good_bad() {
        let config = TpeConfig::new(2, 10);
        let mut tpe = TpeOptimizer::new(config);
        tpe.add_trial(vec![0.5, 0.5], 10.0);
        tpe.add_trial(vec![0.3, 0.7], 5.0);
        tpe.add_trial(vec![0.8, 0.2], 20.0);
        tpe.add_trial(vec![0.2, 0.8], 1.0);

        let (good, bad) = tpe.split_trials();
        assert_eq!(good.len(), 1); // top 25% of 4 = 1
        assert_eq!(bad.len(), 3);
        assert!((good[0].objective - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_tpe_suggest_within_bounds() {
        let config = TpeConfig {
            n_dims: 2,
            max_trials: 10,
            n_initial: 5,
            gamma: 0.25,
            n_candidates: 100,
            seed: Some(123),
        };
        let mut tpe = TpeOptimizer::new(config);
        let bounds = vec![(0.0, 1.0), (0.0, 1.0)];

        // Add initial trials
        tpe.add_trial(vec![0.5, 0.5], 5.0);
        tpe.add_trial(vec![0.3, 0.7], 3.0);
        tpe.add_trial(vec![0.7, 0.3], 8.0);
        tpe.add_trial(vec![0.2, 0.8], 1.0);
        tpe.add_trial(vec![0.1, 0.9], 2.0);
        tpe.add_trial(vec![0.9, 0.1], 15.0);

        let suggestion = tpe.suggest(&bounds);
        assert_eq!(suggestion.len(), 2);
        assert!(
            suggestion[0] >= 0.0 && suggestion[0] <= 1.0,
            "dim 0 out of bounds: {}",
            suggestion[0]
        );
        assert!(
            suggestion[1] >= 0.0 && suggestion[1] <= 1.0,
            "dim 1 out of bounds: {}",
            suggestion[1]
        );
    }

    #[test]
    fn test_tpe_lhs_initial() {
        let config = TpeConfig {
            n_dims: 3,
            max_trials: 10,
            n_initial: 5,
            gamma: 0.25,
            n_candidates: 50,
            seed: Some(42),
        };
        let mut tpe = TpeOptimizer::new(config);
        let bounds = vec![(0.0, 10.0), (-1.0, 1.0), (100.0, 200.0)];

        // First suggestion should be LHS (random within bounds)
        let s = tpe.suggest(&bounds);
        assert_eq!(s.len(), 3);
        assert!(s[0] >= 0.0 && s[0] <= 10.0, "dim 0: {}", s[0]);
        assert!(s[1] >= -1.0 && s[1] <= 1.0, "dim 1: {}", s[1]);
        assert!(s[2] >= 100.0 && s[2] <= 200.0, "dim 2: {}", s[2]);
    }

    #[test]
    fn test_tpe_kde_log_pdf() {
        // KDE of a single point should peak near that point
        let samples = vec![vec![0.5, 0.5]];
        let sample_refs: Vec<&Vec<f64>> = samples.iter().collect();
        let bounds = vec![(0.0, 1.0), (0.0, 1.0)];

        let pdf_at_sample = TpeOptimizer::kde_log_pdf(&sample_refs, &[0.5, 0.5], &bounds);
        let pdf_far = TpeOptimizer::kde_log_pdf(&sample_refs, &[0.0, 0.0], &bounds);

        assert!(
            pdf_at_sample > pdf_far,
            "PDF should be higher near sample point: at={pdf_at_sample}, far={pdf_far}"
        );
    }

    #[test]
    fn test_tpe_kde_log_pdf_multiple_samples() {
        // KDE with multiple samples should produce a reasonable density
        let samples = vec![
            vec![0.2, 0.8],
            vec![0.3, 0.7],
            vec![0.25, 0.75],
        ];
        let sample_refs: Vec<&Vec<f64>> = samples.iter().collect();
        let bounds = vec![(0.0, 1.0), (0.0, 1.0)];

        let pdf_near = TpeOptimizer::kde_log_pdf(&sample_refs, &[0.25, 0.75], &bounds);
        let pdf_far = TpeOptimizer::kde_log_pdf(&sample_refs, &[0.9, 0.1], &bounds);

        assert!(
            pdf_near > pdf_far,
            "PDF should be higher near cluster: near={pdf_near}, far={pdf_far}"
        );
    }

    #[test]
    fn test_tpe_convergence_simple() {
        // TPE should converge toward the minimum of a simple quadratic
        let config = TpeConfig {
            n_dims: 1,
            max_trials: 30,
            n_initial: 5,
            gamma: 0.25,
            n_candidates: 50,
            seed: Some(42),
        };
        let mut tpe = TpeOptimizer::new(config);
        let bounds = vec![(0.0, 10.0)];

        // Objective: (x - 3)^2, minimum at x=3
        for _ in 0..30 {
            let suggestion = tpe.suggest(&bounds);
            let x = suggestion[0];
            let objective = (x - 3.0) * (x - 3.0);
            tpe.add_trial(suggestion, objective);
        }

        // Best trial should be near x=3
        let best = tpe
            .best_trial()
            .expect("should have at least one trial");
        assert!(
            (best.params[0] - 3.0).abs() < 1.5,
            "Best should be near 3.0, got {}",
            best.params[0]
        );
    }

    #[test]
    fn test_tpe_convergence_2d() {
        // 2D Rosenbrock-like: (x - 2)^2 + (y - 3)^2
        let config = TpeConfig {
            n_dims: 2,
            max_trials: 50,
            n_initial: 10,
            gamma: 0.25,
            n_candidates: 100,
            seed: Some(99),
        };
        let mut tpe = TpeOptimizer::new(config);
        let bounds = vec![(0.0, 5.0), (0.0, 5.0)];

        for _ in 0..50 {
            let suggestion = tpe.suggest(&bounds);
            let x = suggestion[0];
            let y = suggestion[1];
            let objective = (x - 2.0) * (x - 2.0) + (y - 3.0) * (y - 3.0);
            tpe.add_trial(suggestion, objective);
        }

        let best = tpe.best_trial().expect("should have trials");
        assert!(
            best.objective < 2.0,
            "Best objective should be < 2.0, got {}",
            best.objective
        );
    }

    #[test]
    fn test_tpe_history_persistence() {
        let config = TpeConfig::new(2, 10);
        let mut tpe = TpeOptimizer::new(config);
        tpe.add_trial(vec![0.5, 0.5], 5.0);
        tpe.add_trial(vec![0.3, 0.7], 3.0);

        let tmp_dir = std::env::temp_dir().join("tpe_test_history");
        std::fs::create_dir_all(&tmp_dir).unwrap();
        let path = tmp_dir.join("history.json");

        tpe.save_history(&path).unwrap();

        let config2 = TpeConfig::new(2, 10);
        let mut tpe2 = TpeOptimizer::new(config2);
        let loaded = tpe2.load_history(&path).unwrap();
        assert_eq!(loaded, 2);
        assert_eq!(tpe2.trials.len(), 2);

        // Verify loaded content
        assert!((tpe2.trials[0].objective - 5.0).abs() < 1e-10);
        assert!((tpe2.trials[1].objective - 3.0).abs() < 1e-10);

        std::fs::remove_dir_all(&tmp_dir).ok();
    }

    #[test]
    fn test_tpe_lhs_coverage() {
        // LHS should produce well-distributed samples across bins
        let config = TpeConfig {
            n_dims: 1,
            max_trials: 10,
            n_initial: 5,
            gamma: 0.25,
            n_candidates: 50,
            seed: Some(7),
        };
        let mut tpe = TpeOptimizer::new(config);
        let bounds = vec![(0.0, 10.0)];

        let samples = tpe.latin_hypercube_sample(&bounds, 10);
        assert_eq!(samples.len(), 10);

        // Each sample should be within bounds
        for s in &samples {
            assert!(s[0] >= 0.0 && s[0] <= 10.0, "out of bounds: {}", s[0]);
        }

        // Check that samples cover different bins (not all clustered)
        let mut bins = [false; 5];
        for s in &samples {
            let bin = ((s[0] / 10.0) * 5.0).floor() as usize;
            let bin = bin.min(4);
            bins[bin] = true;
        }
        let covered = bins.iter().filter(|&&b| b).count();
        assert!(
            covered >= 3,
            "LHS should cover at least 3 of 5 bins, covered {covered}"
        );
    }

    #[test]
    fn test_tpe_empty_bounds() {
        let config = TpeConfig::new(0, 10);
        let mut tpe = TpeOptimizer::new(config);
        let samples = tpe.latin_hypercube_sample(&[], 5);
        assert!(samples.is_empty());
    }

    #[test]
    fn test_tpe_scott_bandwidth() {
        // Identical values should produce floor bandwidth
        let values = vec![5.0, 5.0, 5.0];
        let bw = TpeOptimizer::scott_bandwidth(&values, Some(10.0), 2);
        assert!(
            (bw - 0.1).abs() < 1e-10,
            "Zero-std should use floor bandwidth, got {bw}"
        );

        // Spread values should produce reasonable bandwidth
        let values2 = vec![0.0, 5.0, 10.0];
        let bw2 = TpeOptimizer::scott_bandwidth(&values2, Some(10.0), 2);
        assert!(bw2 > 0.1, "Spread values should have larger bandwidth: {bw2}");
        assert!(bw2 < 10.0, "Bandwidth should be smaller than range: {bw2}");
    }

    #[test]
    fn test_tpe_log_sum_exp() {
        // log(exp(1) + exp(2)) = log(e + e^2) = log(e^2 * (e^-1 + 1)) = 2 + log(1 + e^-1)
        let result = TpeOptimizer::log_sum_exp(&[1.0, 2.0]);
        let expected = (1.0_f64.exp() + 2.0_f64.exp()).ln();
        assert!(
            (result - expected).abs() < 1e-10,
            "log_sum_exp([1, 2]) = {result}, expected {expected}"
        );

        // Edge case: empty
        assert_eq!(TpeOptimizer::log_sum_exp(&[]), f64::NEG_INFINITY);

        // Edge case: single value
        assert!((TpeOptimizer::log_sum_exp(&[5.0]) - 5.0).abs() < 1e-10);
    }

    #[test]
    fn test_tpe_best_trial() {
        let config = TpeConfig::new(1, 10);
        let mut tpe = TpeOptimizer::new(config);
        assert!(tpe.best_trial().is_none());

        tpe.add_trial(vec![1.0], 10.0);
        tpe.add_trial(vec![2.0], 5.0);
        tpe.add_trial(vec![3.0], 8.0);

        let best = tpe.best_trial().unwrap();
        assert!((best.objective - 5.0).abs() < 1e-10);
        assert!((best.params[0] - 2.0).abs() < 1e-10);
    }

    #[test]
    fn test_tpe_n_trials() {
        let config = TpeConfig::new(1, 10);
        let mut tpe = TpeOptimizer::new(config);
        assert_eq!(tpe.n_trials(), 0);

        tpe.add_trial(vec![1.0], 1.0);
        assert_eq!(tpe.n_trials(), 1);

        tpe.add_trial(vec![2.0], 2.0);
        assert_eq!(tpe.n_trials(), 2);
    }

    #[test]
    fn test_tpe_split_preserves_all_trials() {
        let config = TpeConfig {
            n_dims: 1,
            max_trials: 10,
            n_initial: 5,
            gamma: 0.25,
            n_candidates: 50,
            seed: None,
        };
        let mut tpe = TpeOptimizer::new(config);

        for i in 0..8 {
            tpe.add_trial(vec![i as f64], i as f64);
        }

        let (good, bad) = tpe.split_trials();
        assert_eq!(
            good.len() + bad.len(),
            8,
            "Split should preserve all trials"
        );
        // gamma=0.25 of 8 = 2
        assert_eq!(good.len(), 2);
        assert_eq!(bad.len(), 6);
    }
}
