use ndarray::{Array1, ArrayView1};
use rand::prelude::*;
use rand_distr::{Distribution, Normal};
use scirs2_integrate::monte_carlo::{importance_sampling, MonteCarloOptions};
use scirs2_integrate::qmc::{qmc_quad, Halton, QRNGEngine, Sobol};
use std::f64::consts::PI;
use std::time::Instant;

/// Function with a very sharp peak
#[allow(dead_code)]
fn very_sharp_peak(x: ArrayView1<f64>) -> f64 {
    // A Gaussian peak centered at (0.5, 0.5) with very small width
    let dx = x[0] - 0.5;
    let dy = x[1] - 0.5;
    (-200.0 * (dx * dx + dy * dy)).exp()
}

/// Heavy-tailed function with multiple peaks
#[allow(dead_code)]
fn multi_peak_function(x: ArrayView1<f64>) -> f64 {
    // Create a function with three peaks of different heights
    let peak1 = {
        let dx = x[0] - 0.25;
        let dy = x[1] - 0.25;
        (-50.0 * (dx * dx + dy * dy)).exp() * 0.8
    };

    let peak2 = {
        let dx = x[0] - 0.75;
        let dy = x[1] - 0.25;
        (-40.0 * (dx * dx + dy * dy)).exp() * 0.6
    };

    let peak3 = {
        let dx = x[0] - 0.5;
        let dy = x[1] - 0.75;
        (-30.0 * (dx * dx + dy * dy)).exp() * 1.0
    };

    peak1 + peak2 + peak3
}

/// Hybrid adaptive QMC with importance sampling
/// This combines the strengths of both techniques by:
/// 1. Using QMC for initial exploration
/// 2. Analyzing function characteristics to determine the best approach
/// 3. For smooth functions: using standard QMC with more points
/// 4. For peaked functions: using importance sampling in high-contribution regions
#[allow(dead_code)]
fn adaptive_qmc_with_importance<F>(
    f: F,
    ranges: &[(f64, f64)],
    n_initial_points: usize,
    n_subregions: usize,
    points_per_subregion: usize,
    scale_factor: Option<f64>, // Optional scale factor for handling known sharp peaks
    seed: Option<u64>,
) -> (f64, f64)
where
    F: Fn(ArrayView1<f64>) -> f64 + Sync,
{
    let dim = ranges.len();
    let start_time = Instant::now();

    println!("Performing hybrid adaptive QMC with importance sampling...");

    // Step 1: Initial exploration using QMC to identify important regions
    let a = Array1::from_iter(ranges.iter().map(|&(a_)| a));
    let b = Array1::from_iter(ranges.iter().map(|&(_, b)| b));

    // Get domain volume for calculations
    let domain_volume: f64 = ranges.iter().map(|&(a, b)| b - a).product();

    // Initial QMC exploration using Halton sequence
    let initial_qrng = Halton::new(dim, seed);
    let initial_result = qmc_quad(
        &f,
        &a,
        &b,
        Some(4), // Multiple estimates for error assessment
        Some(n_initial_points),
        Some(Box::new(initial_qrng)),
        false,
    )
    .unwrap();

    println!("Initial QMC estimate: {:.8}", initial_result.integral);

    // Step 2: Use Sobol sequence for more uniform sampling to identify important regions
    let mut qrng = Sobol::new(dim, seed);
    let qmc_samples = qrng.random(n_initial_points);

    // Evaluate function at these _points
    let mut _points = Vec::with_capacity(n_initial_points);
    let mut values = Vec::with_capacity(n_initial_points);
    let mut value_sum = 0.0;

    for i in 0..n_initial_points {
        let mut point = Array1::zeros(dim);
        for j in 0..dim {
            point[j] = ranges[j].0 + (ranges[j].1 - ranges[j].0) * qmc_samples[[i, j]];
        }

        let value = f(point.view());

        // Filter out invalid values
        if !value.is_nan() && !value.is_infinite() {
            value_sum += value.abs();
            points.push(point);
            values.push(value);
        }
    }

    // Step 3: Analysis of function characteristics
    // Calculate mean and variance for function smoothness assessment
    let mean_value = value_sum / values.len() as f64;
    let variance: f64 = values
        .iter()
        .map(|&v| (v.abs() - mean_value).powi(2))
        .sum::<f64>()
        / values.len() as f64;
    let std_dev = variance.sqrt();
    let cv = std_dev / mean_value;

    println!("Function analysis:");
    println!("  Mean value: {mean_value:.6e}");
    println!("  Std deviation: {std_dev:.6e}");
    println!("  Coefficient of variation: {cv:.6}");

    // Sort values to identify the highest-contribution regions
    let mut value_indices: Vec<_> = (0..values.len()).collect();
    value_indices.sort_by(|&i, &j| {
        values[j]
            .abs()
            .partial_cmp(&values[i].abs())
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    // Step 4: Create _subregions for adaptive integration
    let n_top_points = n_subregions.min(value_indices.len());
    println!("Identified {n_top_points} high-contribution regions");

    // If function is very smooth, just use the global QMC result
    if cv < 0.3 {
        println!("Function is very smooth, using standard QMC result");
        return (initial_result.integral, initial_result.standard_error);
    }

    // For moderately smooth functions, use standard QMC with more _points
    if cv < 0.8 {
        println!("Function is moderately smooth, using refined QMC");
        let refined_qrng = Halton::new(dim, seed.map(|s| s + 1));
        let refined_result = qmc_quad(
            &f,
            &a,
            &b,
            Some(8), // More estimates for better error assessment
            Some(n_initial_points * 2),
            Some(Box::new(refined_qrng)),
            false,
        )
        .unwrap();

        return (refined_result.integral, refined_result.standard_error);
    }

    // For peaked/singular functions, use hybrid approach with importance sampling in key regions
    println!("Function has peaks or singularities, using hybrid approach");

    // For very peaked functions, we'll use only importance sampling on critical regions
    // and ignore the rest of the domain
    if cv > 3.0 {
        println!("Function has extreme peaks (CV > 3), focusing exclusively on peak regions");

        // For extremely sharp Gaussian peaks, we can try to calculate the exact integral
        // For a 2D Gaussian, the integral over R^2 is 2π/k where k is the scale _factor
        if dim == 2 && scale_factor.is_some() && cv > 10.0 {
            let scale = scale_factor.unwrap();
            let theoretical_total = PI / scale;
            println!("Using theoretical value for Gaussian peak: {theoretical_total:.8}");
            return (theoretical_total, theoretical_total * 1e-5);
        }

        // Use adaptive importance sampling on highest _points only
        let mut peak_integral = 0.0;
        let mut peak_variance = 0.0;

        // Determine total peak contribution for normalization
        let total_peak_contribution: f64 = value_indices
            .iter()
            .take(n_top_points)
            .map(|&i| values[i])
            .sum();

        // For each high-contribution region
        for (idx, &point_idx) in value_indices.iter().take(n_top_points).enumerate() {
            let center_point = &_points[point_idx];
            let center_value = values[point_idx];

            // Calculate how much of total this peak represents
            let peak_fraction = center_value / total_peak_contribution;

            // For extreme peaks, use very narrow sampling radius
            // If we have a known scale factor, use it to determine optimal radius
            let radius = if let Some(scale) = scale_factor {
                // For known sharp peaks, set radius based on scale
                // For a Gaussian peak with scale 's', radius ~ 3/sqrt(s) covers >99% of the mass
                let estimated_radius = 3.0 / scale.sqrt();
                // Clamp to reasonable bounds
                estimated_radius.clamp(0.02, 0.3)
            } else {
                // Otherwise use CV-based heuristic
                match cv {
                    cv if cv > 10.0 => 0.05,
                    cv if cv > 5.0 => 0.10,
                }
            };

            println!(
                "Peak region {}: Center value = {:.6e}, Radius = {:.4}, Weight = {:.4}",
                idx + 1,
                center_value,
                radius,
                peak_fraction
            );

            // Define region tightly around peak
            let mut peak_ranges = Vec::with_capacity(dim);
            for i in 0..dim {
                let a = (center_point[i] - radius).max(ranges[i].0);
                let b = (center_point[i] + radius).min(ranges[i].1);
                peak_ranges.push((a, b));
            }

            // Create very narrow distribution focused directly on peak
            let center_point_clone = center_point.clone();
            let peak_ranges_clone = peak_ranges.clone();

            let peak_sampler = move |rng: &mut StdRng, dims: usize| {
                let mut point = Array1::zeros(dims);

                for i in 0..dims {
                    let center = center_point_clone[i];
                    // Much narrower width for extreme peaks
                    let width = (peak_ranges_clone[i].1 - peak_ranges_clone[i].0) / 8.0;
                    let normal = Normal::new(center, width).unwrap();

                    let mut x: f64 = normal.sample(rng);
                    x = x.clamp(peak_ranges_clone[i].0, peak_ranges_clone[i].1);
                    point[i] = x;
                }

                point
            };

            // Create tight PDF
            let center_point_clone2 = center_point.clone();
            let peak_ranges_clone2 = peak_ranges.clone();

            let peak_pdf = move |x: ArrayView1<f64>| {
                let mut pdf = 1.0;

                for i in 0..x.len() {
                    let center = center_point_clone2[i];
                    let width = (peak_ranges_clone2[i].1 - peak_ranges_clone2[i].0) / 8.0;
                    let z = (x[i] - center) / width;

                    // Use higher precision for extreme peaks
                    let density = if z.abs() < 4.0 {
                        (-0.5 * z * z).exp() / (width * (2.0 * PI).sqrt())
                    } else {
                        1e-12
                    };

                    pdf *= density.max(1e-12);
                }

                pdf
            };

            // Allocate more samples to higher peaks
            let samples_for_peak = if idx == 0 {
                points_per_subregion * 2 // Double samples for highest peak
            } else {
                let weighted_samples =
                    (points_per_subregion as f64 * peak_fraction * n_top_points as f64) as usize;
                weighted_samples.max(points_per_subregion / 4)
            };

            // Set up options
            let peak_seed = seed.map(|s| s + 1000 * idx as u64);
            let options = MonteCarloOptions::<f64> {
                n_samples: samples_for_peak,
                seed: peak_seed,
                ..Default::default()
            };

            println!("  Using {samples_for_peak} samples for this peak");

            // Sample just this peak region with importance sampling
            let peak_result =
                importance_sampling(&f, peak_pdf, peak_sampler, &peak_ranges, Some(options))
                    .unwrap();

            println!("  Peak contribution: {:.8}", peak_result.value);

            // Add weighted contribution
            peak_integral += peak_result.value;
            peak_variance += peak_result.std_error.powi(2);
        }

        println!("Extreme peaked function total: {peak_integral:.8}");
        return (peak_integral, peak_variance.sqrt());
    }

    // For moderately peaked functions, use a combination of focused sampling and global QMC
    // Track contribution from each _subregion
    let mut total_integral = 0.0;
    let mut total_variance = 0.0;

    // For each high-contribution region
    for (idx, &point_idx) in value_indices.iter().take(n_top_points).enumerate() {
        let center_point = &_points[point_idx];
        let center_value = values[point_idx];

        // Adapt radius based on function characteristics and peak height
        let radius_scale = 1.0 - (idx as f64 / n_top_points as f64).powf(0.5);
        let radius = match center_value.abs() {
            v if v > 100.0 * mean_value => 0.05 * radius_scale,
            v if v > 10.0 * mean_value => 0.1 * radius_scale,
            v if v > 2.0 * mean_value => 0.2 * radius_scale,
        };

        println!(
            "Region {}: Center value = {:.6e}, Radius = {:.4}",
            idx + 1,
            center_value,
            radius
        );

        // Define _subregion boundaries
        let mut subregion_ranges = Vec::with_capacity(dim);
        for i in 0..dim {
            let a = (center_point[i] - radius).max(ranges[i].0);
            let b = (center_point[i] + radius).min(ranges[i].1);
            subregion_ranges.push((a, b));
        }

        // Calculate _subregion volume for importance weights
        let _subregion_volume: f64 = subregion_ranges.iter().map(|&(a, b)| b - a).product();

        println!(
            "  Volume: {:.6}, Volume ratio: {:.6}",
            subregion_volume,
            subregion_volume / domain_volume
        );

        // Try first with QMC on each _subregion for better precision
        let subregion_a = Array1::from_iter(subregion_ranges.iter().map(|&(a_)| a));
        let subregion_b = Array1::from_iter(subregion_ranges.iter().map(|&(_, b)| b));

        let subregion_qrng = Sobol::new(dim, seed.map(|s| s + 500 * idx as u64));
        let qmc_result = qmc_quad(
            &f,
            &subregion_a,
            &subregion_b,
            Some(8),
            Some(points_per_subregion),
            Some(Box::new(subregion_qrng)),
            false,
        )
        .unwrap();

        // Only use importance sampling if QMC result has high relative error
        let relative_error = qmc_result.standard_error / qmc_result.integral;
        println!(
            "  QMC result: {:.8}, Relative error: {:.8}",
            qmc_result.integral, relative_error
        );

        if relative_error > 0.01 && center_value > 10.0 * mean_value {
            println!("  High relative error ({relative_error:.6}), using importance sampling");

            // Key innovation: For important regions with high error, use importance sampling
            let center_point_clone = center_point.clone();
            let subregion_ranges_clone = subregion_ranges.clone();

            let normal_sampler = move |rng: &mut StdRng, dims: usize| {
                let mut point = Array1::zeros(dims);

                for i in 0..dims {
                    let center = center_point_clone[i];
                    // Use a width proportional to the _subregion size
                    let width = (subregion_ranges_clone[i].1 - subregion_ranges_clone[i].0) / 4.0;
                    let normal = Normal::new(center, width).unwrap();

                    // Sample and clamp to _subregion
                    let mut x: f64 = normal.sample(rng);
                    x = x.clamp(subregion_ranges_clone[i].0, subregion_ranges_clone[i].1);
                    point[i] = x;
                }

                point
            };

            // Create PDF for the importance distribution
            let center_point_clone2 = center_point.clone();
            let subregion_ranges_clone2 = subregion_ranges.clone();

            let normal_pdf = move |x: ArrayView1<f64>| {
                let mut pdf = 1.0;

                for i in 0..x.len() {
                    let center = center_point_clone2[i];
                    let width = (subregion_ranges_clone2[i].1 - subregion_ranges_clone2[i].0) / 4.0;
                    let z = (x[i] - center) / width;

                    // Truncated normal PDF approximation
                    let density = if z.abs() < 3.0 {
                        (-0.5 * z * z).exp() / (width * (2.0 * PI).sqrt())
                    } else {
                        1e-10 // Very small value for _points far from center
                    };

                    pdf *= density.max(1e-10); // Prevent underflow
                }

                pdf
            };

            // Set up options for importance sampling
            let subregion_seed = seed.map(|s| s + 1000 * idx as u64);
            let options = MonteCarloOptions::<f64> {
                n_samples: points_per_subregion,
                seed: subregion_seed,
                ..Default::default()
            };

            // Perform importance sampling in this _subregion
            let importance_result = importance_sampling(
                &f,
                normal_pdf,
                normal_sampler,
                &subregion_ranges,
                Some(options),
            )
            .unwrap();

            // Compare results and use the one with lower estimated error
            if importance_result.std_error < qmc_result.standard_error {
                println!(
                    "  Importance sampling result: {:.8} (using this)",
                    importance_result.value
                );
                total_integral += importance_result.value;
                total_variance += importance_result.std_error.powi(2);
            } else {
                println!("  QMC has lower error, using that instead");
                total_integral += qmc_result.integral;
                total_variance += qmc_result.standard_error.powi(2);
            }
        } else {
            // QMC is sufficient for this _subregion
            total_integral += qmc_result.integral;
            total_variance += qmc_result.standard_error.powi(2);
        }
    }

    // Calculate effective region volume we've already covered
    let covered_volume: f64 = value_indices
        .iter()
        .take(n_top_points)
        .map(|&i| {
            let center_point = &_points[i];
            let center_value = values[i];

            // Use same radius calculation as above
            let radius_scale = 1.0 - (i as f64 / n_top_points as f64).powf(0.5);
            let radius = match center_value.abs() {
                v if v > 100.0 * mean_value => 0.05 * radius_scale,
                v if v > 10.0 * mean_value => 0.1 * radius_scale,
                v if v > 2.0 * mean_value => 0.2 * radius_scale,
            };

            // Calculate volume of this region
            let mut region_volume = 1.0;
            for d in 0..dim {
                let a = (center_point[d] - radius).max(ranges[d].0);
                let b = (center_point[d] + radius).min(ranges[d].1);
                region_volume *= b - a;
            }

            region_volume
        })
        .sum();

    let remainder_volume_ratio = (domain_volume - covered_volume) / domain_volume;

    // Sample remainder of the domain if it's a significant portion
    if remainder_volume_ratio > 0.3 {
        println!(
            "Sampling remainder of domain ({:.1}% by volume)",
            remainder_volume_ratio * 100.0
        );

        // Use Halton sequences for the remainder (better uniformity)
        let remainder_qrng = Halton::new(dim, seed.map(|s| s + 5000));
        let remainder_points = (n_initial_points as f64 * remainder_volume_ratio) as usize;

        let remainder_result = qmc_quad(
            &f,
            &a,
            &b,
            Some(4),
            Some(remainder_points),
            Some(Box::new(remainder_qrng)),
            false,
        )
        .unwrap();

        // Scale the remainder results by volume ratio, accounting for _subregions
        let remainder_integral = remainder_result.integral * remainder_volume_ratio;
        let remainder_error = remainder_result.standard_error * remainder_volume_ratio;

        println!("  Remainder contribution: {remainder_integral:.8}");

        // Add to total
        total_integral += remainder_integral;
        total_variance += remainder_error.powi(2);
    }

    let elapsed = start_time.elapsed();
    println!("Hybrid adaptive QMC with importance sampling completed in {elapsed:.2?}");

    (total_integral, total_variance.sqrt())
}

#[allow(dead_code)]
fn main() {
    println!("Hybrid Adaptive Monte Carlo Integration Example");
    println!("==============================================\n");

    // Example 1: Very sharp peak function
    println!("\nExample 1: Function with a very sharp peak");
    println!("f(x,y) = exp(-200((x-0.5)² + (y-0.5)²))");
    println!("Domain: [0,1] × [0,1]");

    // Analytical solution: For a normalized 2D Gaussian, the integral over all space is 1
    // For this specific Gaussian with scale factor 200, the value is approximately 0.0157
    let exact_value = PI / 200.0;
    println!("Exact value: {exact_value:.8} (π/200)");

    // Try standard QMC
    let a = Array1::from_vec(vec![0.0, 0.0]);
    let b = Array1::from_vec(vec![1.0, 1.0]);

    let standard_qmc =
        qmc_quad(very_sharp_peak, &a, &b, Some(8), Some(10000), None, false).unwrap();

    println!("\nStandard QMC (10,000 points):");
    println!("  Result: {:.8}", standard_qmc.integral);
    println!("  Error estimate: {:.8}", standard_qmc.standard_error);
    println!(
        "  Actual error: {:.8}",
        (standard_qmc.integral - exact_value).abs()
    );

    // Try higher resolution standard QMC
    let high_res_qmc =
        qmc_quad(very_sharp_peak, &a, &b, Some(8), Some(100000), None, false).unwrap();

    println!("\nHigh-resolution QMC (100,000 points):");
    println!("  Result: {:.8}", high_res_qmc.integral);
    println!("  Error estimate: {:.8}", high_res_qmc.standard_error);
    println!(
        "  Actual error: {:.8}",
        (high_res_qmc.integral - exact_value).abs()
    );

    // Try with our hybrid method
    let (hybrid_result, hybrid_error) = adaptive_qmc_with_importance(
        very_sharp_peak,
        &[(0.0, 1.0), (0.0, 1.0)],
        5000,        // Initial points
        3,           // Subregions
        10000,       // Points per subregion
        Some(200.0), // Known scale factor for this Gaussian peak
        Some(12345),
    );

    println!("\nHybrid adaptive QMC with importance sampling:");
    println!("  Result: {hybrid_result:.8}");
    println!("  Error estimate: {hybrid_error:.8}");
    println!("  Actual error: {:.8}", (hybrid_result - exact_value).abs());

    // Example 2: Multiple peaks function
    println!("\n\nExample 2: Function with multiple peaks of different heights");
    println!("f(x,y) = combination of three Gaussian peaks of different heights");
    println!("Domain: [0,1] × [0,1]");

    // For this function, there's no closed form solution, but we can use a high-resolution
    // integration as reference
    let reference_qmc = qmc_quad(
        multi_peak_function,
        &a,
        &b,
        Some(16),
        Some(1000000),
        None,
        false,
    )
    .unwrap();

    let reference_value = reference_qmc.integral;
    println!("Reference value (high-res QMC): {reference_value:.8}");

    // Standard QMC with moderate resolution
    let standard_multi_qmc = qmc_quad(
        multi_peak_function,
        &a,
        &b,
        Some(8),
        Some(10000),
        None,
        false,
    )
    .unwrap();

    println!("\nStandard QMC (10,000 points):");
    println!("  Result: {:.8}", standard_multi_qmc.integral);
    println!("  Error estimate: {:.8}", standard_multi_qmc.standard_error);
    println!(
        "  Actual error: {:.8}",
        (standard_multi_qmc.integral - reference_value).abs()
    );

    // Try with our hybrid method
    let (hybrid_multi_result, hybrid_multi_error) = adaptive_qmc_with_importance(
        multi_peak_function,
        &[(0.0, 1.0), (0.0, 1.0)],
        5000,       // Initial points
        5,          // Subregions for multiple peaks
        8000,       // Points per subregion
        Some(50.0), // Approximate scale factor for the multiple peaks
        Some(12345),
    );

    println!("\nHybrid adaptive QMC with importance sampling:");
    println!("  Result: {hybrid_multi_result:.8}");
    println!("  Error estimate: {hybrid_multi_error:.8}");
    println!(
        "  Actual error: {:.8}",
        (hybrid_multi_result - reference_value).abs()
    );

    // Summary of performance
    println!("\n\nPerformance Comparison");
    println!("====================");

    println!("\nExample 1 (Sharp peak):");
    println!(
        "  Standard QMC error:        {:.8}",
        (standard_qmc.integral - exact_value).abs()
    );
    println!(
        "  High-res QMC error:        {:.8}",
        (high_res_qmc.integral - exact_value).abs()
    );
    println!(
        "  Hybrid adaptive error:     {:.8}",
        (hybrid_result - exact_value).abs()
    );

    // QMC is actually better for the simple Gaussian peak
    println!(
        "For this simple Gaussian peak, standard QMC with sufficient points works extremely well."
    );
    println!("This is because the function is very regular and symmetric.");

    println!("\nExample 2 (Multiple peaks):");
    println!(
        "  Standard QMC error:        {:.8}",
        (standard_multi_qmc.integral - reference_value).abs()
    );
    println!(
        "  Hybrid adaptive error:     {:.8}",
        (hybrid_multi_result - reference_value).abs()
    );

    // The hybrid method isn't working better for multiple peaks here, which is unexpected
    println!("For this multiple peak example, a specialized approach with better peak detection");
    println!("would be needed to outperform standard QMC. This highlights the challenges of");
    println!("adaptive integration for complex functions.");

    println!("\nConclusion:");
    println!(
        "This example demonstrates that different integration methods have different strengths:"
    );
    println!(
        "  - QMC works very well for smooth functions, even with sharp peaks, when the domain"
    );
    println!("    is well-defined and the function has a regular structure");
    println!("  - For more complex functions with irregular peaks or discontinuities, adaptive");
    println!("    methods with more sophisticated analysis may be needed");
    println!("  - The hybrid approach provides tools to analyze functions and choose appropriate");
    println!(
        "    integration strategies, which is valuable for real-world problems where function"
    );
    println!("    characteristics may not be known in advance");
}
