#![allow(dead_code)]

use clap::Parser;
use faer::traits::ComplexField;
use faer::Mat;
use num_traits::{Float, FromPrimitive, ToPrimitive};
use rand::{rngs::StdRng, seq::SliceRandom, Rng, SeedableRng};
use rand_distr::StandardNormal;
use rustc_hash::FxHashSet;
use thousands::*;

////////////
// Consts //
////////////

/// Default number of samples
pub const DEFAULT_N_SAMPLES: usize = 150_000;
/// Default number for the querying
pub const DEFAULT_N_QUERY: usize = DEFAULT_N_SAMPLES / 10;
/// Default dimensionality -> typical for single cell
pub const DEFAULT_DIM: usize = 32;
/// Number of default clusters
pub const DEFAULT_N_CLUSTERS: usize = 25;
/// Default number of neighbours
pub const DEFAULT_K: usize = 15;
/// Default random seed
pub const DEFAULT_SEED: u64 = 42;
/// Default distance metric
pub const DEFAULT_DISTANCE: &str = "euclidean";
/// Correlation strength for
pub const DEFAULT_COR_STRENGTH: f64 = 0.5;
/// Default data type
pub const DEFAULT_DATA: &str = "gaussian";
/// Default intrinsic dimensions
pub const DEFAULT_INTRINSIC_DIM: usize = 16;
/// Default spectral decay
pub const DEFAULT_SPECTRAL_DECAY: f64 = 1.5;

////////////
// Parser //
////////////

/// Parsing structure
///
/// ### Fields
///
/// * `n_samples` - Number of samples/samples
/// * `dim` - Number of dimensions to use
/// * `n_clusters` - Number of clusters in the data
/// * `k` - Number of neighbours to search
/// * `seed` - Random seed for reproducibility
/// * `distance` - The distance to use. One of `"euclidean"` or `"cosine"`.
/// * `data` - The data to use. One of `"gaussian"` or `"correlated"`.
#[derive(Parser, Clone)]
pub struct Cli {
    #[arg(long, default_value_t = DEFAULT_N_SAMPLES)]
    pub n_samples: usize,

    #[arg(long, default_value_t = DEFAULT_DIM)]
    pub dim: usize,

    #[arg(long, default_value_t = DEFAULT_N_CLUSTERS)]
    pub n_clusters: usize,

    #[arg(long, default_value_t = DEFAULT_K)]
    pub k: usize,

    #[arg(long, default_value_t = DEFAULT_SEED)]
    pub seed: u64,

    #[arg(long, default_value = DEFAULT_DISTANCE)]
    pub distance: String,

    #[arg(long, default_value = DEFAULT_DATA)]
    pub data: String,

    #[arg(long, default_value_t = DEFAULT_INTRINSIC_DIM)]
    pub intrinsic_dim: usize,

    #[arg(long, default_value_t = DEFAULT_SPECTRAL_DECAY)]
    pub spectral_decay: f64,
}

//////////
// Data //
//////////

/// Enum defining the synthetic data
#[derive(Default)]
pub enum SyntheticData {
    /// Default Gaussian noise cluster
    #[default]
    GaussianNoise,
    /// Correlated structure
    Correlated,
    /// LowRank type of data
    LowRank,
    /// Data type to stress quantisation methods
    QuantisationStress,
}

/// Helper function to parse the data type
///
/// ### Params
///
/// * `s` - The string to parse
///
/// ### Returns
///
/// `Option<SyntheticData>`
pub fn parse_data(s: &str) -> Option<SyntheticData> {
    match s.to_lowercase().as_str() {
        "gaussian" => Some(SyntheticData::GaussianNoise),
        "correlated" => Some(SyntheticData::Correlated),
        "lowrank" => Some(SyntheticData::LowRank),
        "quantisation" | "quantization" => Some(SyntheticData::QuantisationStress),
        _ => None,
    }
}

/// Generate synthetic single-cell-like data with cluster structure
///
/// Creates data with multiple Gaussian clusters to simulate clusters, cell
/// types in the data
///
/// ### Params
///
/// * `n_samples` - Number of samples (samples)
/// * `dim` - Embedding dimensionality
/// * `n_clusters` - Number of distinct clusters
/// * `seed` - Random seed for reproducibility
///
/// ### Returns
///
/// Matrix of shape (n_samples, dim)
pub fn generate_clustered_data<T>(
    n_samples: usize,
    dim: usize,
    n_clusters: usize,
    seed: u64,
) -> (Mat<T>, Vec<usize>)
where
    T: Float + FromPrimitive + ComplexField,
{
    let mut rng = StdRng::seed_from_u64(seed);
    let mut data = Mat::<T>::zeros(n_samples, dim);

    // variable cluster sizes and std deviations
    let mut centres = Vec::with_capacity(n_clusters);
    let mut cluster_stds = Vec::new();

    for _ in 0..n_clusters {
        let centre: Vec<f64> = (0..dim).map(|_| rng.random_range(-7.5..7.5)).collect();
        centres.push(centre);
        cluster_stds.push(rng.random_range(0.5..2.5));
    }

    // assign samples with variable cluster sizes
    // with some clusters bigger than others
    let mut cluster_assignments = Vec::new();
    for cluster_idx in 0..n_clusters {
        let weight = rng.random_range(0.5..2.5);
        let n_in_cluster = ((n_samples as f64 * weight) / (n_clusters as f64 * 1.25)) as usize;
        cluster_assignments.extend(vec![cluster_idx; n_in_cluster]);
    }

    // fill remaining
    while cluster_assignments.len() < n_samples {
        cluster_assignments.push(rng.random_range(0..n_clusters));
    }
    cluster_assignments.shuffle(&mut rng);
    cluster_assignments.truncate(n_samples);

    // generate with variable noise
    for (i, &cluster_idx) in cluster_assignments.iter().enumerate() {
        let centre = &centres[cluster_idx];
        let std = cluster_stds[cluster_idx];

        for j in 0..dim {
            let u1: f64 = rng.random();
            let u2: f64 = rng.random();
            let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
            data[(i, j)] = T::from_f64(centre[j] + noise * std).unwrap();
        }
    }

    (data, cluster_assignments)
}

/// Generate synthetic single-cell-like data with cluster structure and
/// correlated dimensions
///
/// Creates well-separated clusters with subspace structure plus inter-dimension
/// correlations that OPQ can exploit.
///
/// ### Params
///
/// * `n_samples` - Number of samples (samples)
/// * `dim` - Embedding dimensionality
/// * `n_clusters` - Number of distinct clusters
/// * `correlation_strength` - How strongly correlated dims depend on source
///   dims (0.0-1.0)
/// * `seed` - Random seed for reproducibility
///
/// ### Returns
///
/// Matrix of shape (n_samples, dim)
pub fn generate_clustered_data_high_dim<T>(
    n_samples: usize,
    dim: usize,
    n_clusters: usize,
    correlation_strength: f64,
    seed: u64,
) -> (Mat<T>, Vec<usize>)
where
    T: Float + FromPrimitive + ComplexField,
{
    let mut rng = StdRng::seed_from_u64(seed);
    let mut data = Mat::<T>::zeros(n_samples, dim);

    let scale = (dim as f64).sqrt() * 2.0;

    // generate well-separated centres
    let mut centres = Vec::with_capacity(n_clusters);
    let min_separation = scale * 0.8;

    for _ in 0..n_clusters {
        let centre = loop {
            let candidate: Vec<f64> = (0..dim).map(|_| rng.random_range(-scale..scale)).collect();
            let too_close = centres.iter().any(|existing: &Vec<f64>| {
                let dist_sq: f64 = candidate
                    .iter()
                    .zip(existing.iter())
                    .map(|(a, b)| (a - b).powi(2))
                    .sum();
                dist_sq < min_separation.powi(2)
            });
            if !too_close {
                break candidate;
            }
        };
        centres.push(centre);
    }

    // subspace structure
    let active_dims_per_cluster = (dim / 2).max(3);
    let cluster_active_dims: Vec<Vec<usize>> = centres
        .iter()
        .map(|_| {
            let mut dims: Vec<usize> = (0..dim).collect();
            dims.shuffle(&mut rng);
            dims.truncate(active_dims_per_cluster);
            dims
        })
        .collect();

    let mut cluster_stds = Vec::new();
    for _ in 0..n_clusters {
        cluster_stds.push(rng.random_range(0.3..1.0) * scale / 10.0);
    }

    // assign samples
    let mut cluster_assignments = Vec::new();
    for cluster_idx in 0..n_clusters {
        let weight = rng.random_range(0.5..2.5);
        let n_in_cluster = ((n_samples as f64 * weight) / (n_clusters as f64 * 1.25)) as usize;
        cluster_assignments.extend(vec![cluster_idx; n_in_cluster]);
    }
    while cluster_assignments.len() < n_samples {
        cluster_assignments.push(rng.random_range(0..n_clusters));
    }
    cluster_assignments.shuffle(&mut rng);
    cluster_assignments.truncate(n_samples);

    // generate base samples
    for (i, &cluster_idx) in cluster_assignments.iter().enumerate() {
        let centre = &centres[cluster_idx];
        let std = cluster_stds[cluster_idx];
        let active_dims = &cluster_active_dims[cluster_idx];

        for j in 0..dim {
            let noise_scale = if active_dims.contains(&j) {
                std
            } else {
                std * 0.1
            };
            let u1: f64 = rng.random();
            let u2: f64 = rng.random();
            let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
            data[(i, j)] = T::from_f64(centre[j] + noise * noise_scale).unwrap();
        }
    }

    // add correlation structure: groups of dimensions that are linear
    // combinations of "source" dimensions plus noise
    let n_correlation_groups = dim / 8;
    let dims_per_group = 4;
    let noise_weight = 1.0 - correlation_strength;

    for group in 0..n_correlation_groups {
        let source_dim = group * 8;
        if source_dim >= dim {
            break;
        }

        // generate random mixing coefficients for this group
        let coeffs: Vec<f64> = (0..dims_per_group)
            .map(|_| rng.random_range(-2.0..2.0))
            .collect();

        for target_offset in 1..=dims_per_group {
            let target_dim = source_dim + target_offset;
            if target_dim >= dim {
                break;
            }

            for i in 0..n_samples {
                let source_val = data[(i, source_dim)].to_f64().unwrap();
                let original_val = data[(i, target_dim)].to_f64().unwrap();

                // correlated value = weighted sum of source + original noise
                let correlated = source_val * coeffs[target_offset - 1] * correlation_strength
                    + original_val * noise_weight;

                data[(i, target_dim)] = T::from_f64(correlated).unwrap();
            }
        }
    }

    (data, cluster_assignments)
}

/// Generate manifold-based data
///
/// Creates high-dimensional data that actually lives in a low-dimensional
/// subspace with rotated cluster structure.
///
/// ### Params
///
/// * `n_samples` - Number of samples
/// * `embedding_dim` - Full dimensionality (e.g., 128, 256)
/// * `intrinsic_dim` - True dimensionality of data (e.g., 16, 32)
/// * `n_clusters` - Number of clusters
/// * `seed` - Random seed
///
/// ### Returns
///
/// Matrix of shape (n_samples, embedding_dim)
pub fn generate_low_rank_rotated_data<T>(
    n_samples: usize,
    embedding_dim: usize,
    intrinsic_dim: usize,
    n_clusters: usize,
    seed: u64,
) -> (Mat<T>, Vec<usize>)
where
    T: Float + FromPrimitive + ComplexField,
{
    assert!(
        intrinsic_dim <= embedding_dim,
        "Intrinsic dim must be <= embedding dim"
    );

    let mut rng = StdRng::seed_from_u64(seed);

    // generate well-separated clusters in low-dimensional space
    let cluster_separation = (intrinsic_dim as f64).sqrt() * 3.0;
    let mut centres_low_dim = Vec::with_capacity(n_clusters);

    for _ in 0..n_clusters {
        let centre = loop {
            let candidate: Vec<f64> = (0..intrinsic_dim)
                .map(|_| rng.random_range(-cluster_separation..cluster_separation))
                .collect();

            let too_close = centres_low_dim.iter().any(|existing: &Vec<f64>| {
                let dist_sq: f64 = candidate
                    .iter()
                    .zip(existing.iter())
                    .map(|(a, b)| (a - b).powi(2))
                    .sum();
                dist_sq < (cluster_separation * 0.5).powi(2)
            });

            if !too_close {
                break candidate;
            }
        };
        centres_low_dim.push(centre);
    }

    // assign samples to clusters
    let mut cluster_assignments = Vec::new();
    for cluster_idx in 0..n_clusters {
        let n_in_cluster = n_samples / n_clusters;
        cluster_assignments.extend(vec![cluster_idx; n_in_cluster]);
    }
    while cluster_assignments.len() < n_samples {
        cluster_assignments.push(rng.random_range(0..n_clusters));
    }
    cluster_assignments.shuffle(&mut rng);

    // generate samples in low-dimensional space
    let mut data_low_dim = Mat::<T>::zeros(n_samples, intrinsic_dim);
    let cluster_std = 0.3;

    for (i, &cluster_idx) in cluster_assignments.iter().enumerate() {
        let centre = &centres_low_dim[cluster_idx];
        for j in 0..intrinsic_dim {
            let u1: f64 = rng.random();
            let u2: f64 = rng.random();
            let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
            data_low_dim[(i, j)] = T::from_f64(centre[j] + noise * cluster_std).unwrap();
        }
    }

    // create random rotation matrix to embed into high-dimensional space
    let mut rotation = Mat::<T>::zeros(intrinsic_dim, embedding_dim);

    // fill with random Gaussian values
    for i in 0..intrinsic_dim {
        for j in 0..embedding_dim {
            let val: f64 = rng.sample(StandardNormal);
            rotation[(i, j)] = T::from_f64(val).unwrap();
        }
    }

    // orthonormalise columns via Gram-Schmidt
    for col in 0..embedding_dim.min(intrinsic_dim) {
        // orthogonalise against previous columns
        for prev_col in 0..col {
            let mut dot = T::zero();
            for row in 0..intrinsic_dim {
                dot = dot + rotation[(row, col)] * rotation[(row, prev_col)];
            }
            for row in 0..intrinsic_dim {
                rotation[(row, col)] = rotation[(row, col)] - dot * rotation[(row, prev_col)];
            }
        }

        let mut norm_sq = T::zero();
        for row in 0..intrinsic_dim {
            norm_sq = norm_sq + rotation[(row, col)] * rotation[(row, col)];
        }
        let norm = norm_sq.sqrt();
        if norm > T::epsilon() {
            for row in 0..intrinsic_dim {
                rotation[(row, col)] = rotation[(row, col)] / norm;
            }
        }
    }

    // project to high-dimensional space: data_high = data_low * rotation
    let mut data_high_dim = Mat::<T>::zeros(n_samples, embedding_dim);
    for i in 0..n_samples {
        for j in 0..embedding_dim {
            let mut sum = T::zero();
            for k in 0..intrinsic_dim {
                sum = sum + data_low_dim[(i, k)] * rotation[(k, j)];
            }
            data_high_dim[(i, j)] = sum;
        }
    }

    // Add small amount of isotropic noise in the full space
    let noise_std = 0.01;
    for i in 0..n_samples {
        for j in 0..embedding_dim {
            let u1: f64 = rng.random();
            let u2: f64 = rng.random();
            let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
            data_high_dim[(i, j)] = data_high_dim[(i, j)] + T::from_f64(noise * noise_std).unwrap();
        }
    }

    (data_high_dim, cluster_assignments)
}

/// Random orthogonal matrix via Gram-Schmidt on Gaussian random matrix.
///
/// ### Params
///
/// * `dim` - Dimensions
/// * `rng` - StdRng for the random sampling
///
/// ### Returns
///
/// Orthogonal matrix
fn random_orthogonal_matrix<T>(dim: usize, rng: &mut StdRng) -> Mat<T>
where
    T: Float + FromPrimitive + ComplexField,
{
    let mut q = Mat::<T>::zeros(dim, dim);

    for i in 0..dim {
        for j in 0..dim {
            let val: f64 = rng.sample(StandardNormal);
            q[(i, j)] = T::from_f64(val).unwrap();
        }
    }

    for col in 0..dim {
        for prev in 0..col {
            let mut dot = T::zero();
            for row in 0..dim {
                dot = dot + q[(row, col)] * q[(row, prev)];
            }
            for row in 0..dim {
                q[(row, col)] = q[(row, col)] - dot * q[(row, prev)];
            }
        }
        let mut norm_sq = T::zero();
        for row in 0..dim {
            norm_sq = norm_sq + q[(row, col)] * q[(row, col)];
        }
        let norm = norm_sq.sqrt();
        if norm > T::epsilon() {
            for row in 0..dim {
                q[(row, col)] = q[(row, col)] / norm;
            }
        }
    }

    q
}

/// Generate combined quantisation stress test data
///
/// Creates high-dimensional data combining power-law eigenvalue spectrum with
/// norm-stratified clusters, randomly rotated so the important subspace doesn't
/// align with coordinate axes. Clusters are paired to share directions at
/// different radii, stressing sign binarisation (norm-blind), axis-aligned PQ
/// (mixes signal and noise partitions), and low-bit methods generally.
///
/// ### Params
///
/// * `n_samples` - Number of samples
/// * `dim` - Embedding dimensionality (e.g., 128, 256)
/// * `n_clusters` - Number of distinct clusters
/// * `spectral_decay` - Exponent for power-law variance decay (1.5 = moderate,
///   2.0 = aggressive)
/// * `seed` - Random seed for reproducibility
///
/// ### Returns
///
/// Matrix of shape (n_samples, dim) and cluster assignments
pub fn generate_quantisation_stress<T>(
    n_samples: usize,
    dim: usize,
    n_clusters: usize,
    spectral_decay: f64,
    seed: u64,
) -> (Mat<T>, Vec<usize>)
where
    T: Float + FromPrimitive + ComplexField,
{
    let mut rng = StdRng::seed_from_u64(seed);

    // power-law eigenvalues: variance concentrated in first few dims
    let eigenvalues: Vec<f64> = (0..dim)
        .map(|i| 1.0 / ((i + 1) as f64).powf(spectral_decay))
        .collect();

    // generate unit directions, then assign radii to create norm stratification
    // half the clusters share directions with a partner at a different radius
    let n_directions = n_clusters.div_ceil(2);
    let radii = [2.0, 8.0, 20.0];

    let mut direction_vecs = Vec::with_capacity(n_directions);
    for _ in 0..n_directions {
        let raw: Vec<f64> = (0..dim)
            .map(|_| rng.sample::<f64, _>(StandardNormal))
            .collect();
        let norm: f64 = raw.iter().map(|x| x * x).sum::<f64>().sqrt();
        direction_vecs.push(raw.into_iter().map(|x| x / norm).collect::<Vec<f64>>());
    }

    // centres: direction * radius, scaled per-dim by eigenvalue
    let mut centres = Vec::with_capacity(n_clusters);
    for cidx in 0..n_clusters {
        let dir = &direction_vecs[cidx % n_directions];
        let radius = radii[cidx % radii.len()];
        let centre: Vec<f64> = (0..dim)
            .map(|d| dir[d] * radius * eigenvalues[d].sqrt())
            .collect();
        centres.push(centre);
    }

    // variable cluster sizes
    let mut assignments = Vec::new();
    for idx in 0..n_clusters {
        let weight = rng.random_range(0.5..2.5);
        let n = ((n_samples as f64 * weight) / (n_clusters as f64 * 1.25)) as usize;
        assignments.extend(vec![idx; n]);
    }
    while assignments.len() < n_samples {
        assignments.push(rng.random_range(0..n_clusters));
    }
    assignments.shuffle(&mut rng);
    assignments.truncate(n_samples);

    // samples with dimension-dependent variance, cluster std proportional to radius
    let mut data = Mat::<T>::zeros(n_samples, dim);
    for (i, &cidx) in assignments.iter().enumerate() {
        let centre = &centres[cidx];
        let radius: f64 = centre.iter().map(|x| x * x).sum::<f64>().sqrt().max(1.0);
        let base_std = radius * 0.06;

        for j in 0..dim {
            let u1: f64 = rng.random();
            let u2: f64 = rng.random();
            let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
            let std = base_std * eigenvalues[j].sqrt();
            data[(i, j)] = T::from_f64(centre[j] + noise * std).unwrap();
        }
    }

    // random rotation so nothing is axis-aligned
    let rotation = random_orthogonal_matrix::<T>(dim, &mut rng);
    let mut rotated = Mat::<T>::zeros(n_samples, dim);
    for i in 0..n_samples {
        for j in 0..dim {
            let mut sum = T::zero();
            for k in 0..dim {
                sum = sum + data[(i, k)] * rotation[(k, j)];
            }
            rotated[(i, j)] = sum;
        }
    }

    (rotated, assignments)
}

////////////////////////////
// Main dispatch function //
////////////////////////////

/// Wrapper function to generate the synthetic data to use
///
/// ### Params
///
/// * `cli` - The Cli structure with the data
///
/// ### Returns
///
/// `(syn data, cluster assignments)`
pub fn generate_data(cli: &Cli) -> (Mat<f32>, Vec<usize>) {
    let data_type = parse_data(&cli.data).unwrap_or_default();
    let res: (Mat<f32>, Vec<usize>) = match data_type {
        SyntheticData::GaussianNoise => {
            println!(">>> Using simple Gaussian cluster data. <<<");
            generate_clustered_data(cli.n_samples, cli.dim, cli.n_clusters, cli.seed)
        }
        SyntheticData::Correlated => {
            println!(">>> Using data with subspace structure and correlated features. <<<");
            generate_clustered_data_high_dim(
                cli.n_samples,
                cli.dim,
                cli.n_clusters,
                DEFAULT_COR_STRENGTH,
                cli.seed,
            )
        }
        SyntheticData::LowRank => {
            println!(">>> Using data that simulating manifold hypothesis. <<<");
            generate_low_rank_rotated_data(
                cli.n_samples,
                cli.dim,
                cli.intrinsic_dim,
                cli.n_clusters,
                cli.seed,
            )
        }
        SyntheticData::QuantisationStress => {
            println!(">>> Using quantisation stress test data. <<<");
            generate_quantisation_stress(
                cli.n_samples,
                cli.dim,
                cli.n_clusters,
                cli.spectral_decay,
                cli.seed,
            )
        }
    };

    res
}

/// Randomly subsample a matrix and add Gaussian noise
///
/// ### Params
///
/// * `data` - The input matrix to subsample
/// * `n_samples` - Number of samples to draw
/// * `noise_scale` - Standard deviation of noise to add
/// * `seed` - Random seed for reproducibility
///
/// ### Returns
///
/// Matrix of shape (n_samples, dim) with noise added
pub fn subsample_with_noise<T>(data: &Mat<T>, n_samples: usize, seed: u64) -> Mat<T>
where
    T: Float + FromPrimitive + ComplexField,
{
    let mut rng = StdRng::seed_from_u64(seed + 1000);
    let (n_rows, n_cols) = data.shape();

    let mut indices: Vec<usize> = (0..n_rows).collect();
    indices.shuffle(&mut rng);
    indices.truncate(n_samples.min(n_rows));

    let mut result = Mat::<T>::zeros(n_samples.min(n_rows), n_cols);

    for (i, &row_idx) in indices.iter().enumerate() {
        for j in 0..n_cols {
            let u1: f64 = rng.random();
            let u2: f64 = rng.random();
            let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
            let noised_value = data[(row_idx, j)].to_f64().unwrap() + noise * 0.05;
            result[(i, j)] = T::from_f64(noised_value).unwrap();
        }
    }

    result
}

////////////////
// Benchmarks //
////////////////

////////////////
// Structures //
////////////////

/// BenchmarkResult
pub struct BenchmarkResultSize {
    /// Name of the method
    pub method: String,
    /// The build time of the index in ms
    pub build_time_ms: f64,
    /// The query time of the index in ms
    pub query_time_ms: f64,
    ///  Total time the index build & query takes in ms
    pub total_time_ms: f64,
    /// Recall@k neighbours against ground truth. Overlap in top k neighbours
    /// for given k
    pub recall_at_k: f64,
    /// Mean distance ratio
    pub mean_dist_rat: f64,
    /// Size of the index
    pub index_size_mb: f64,
}

/// BenchmarkResultPurity - includes cluster purity metric
pub struct BenchmarkResultPurity {
    /// Name of the method
    pub method: String,
    /// The build time of the index in ms
    pub build_time_ms: f64,
    /// The query time of the index in ms
    pub query_time_ms: f64,
    ///  Total time the index build & query takes in ms
    pub total_time_ms: f64,
    /// Recall@k neighbours against ground truth. Overlap in top k neighbours
    /// for given k
    pub recall_at_k: f64,
    /// Fraction of neighbors from same cluster
    pub cluster_purity: f64,
    /// Index size in MB
    pub index_size_mb: f64,
}

/////////////
// Helpers //
/////////////

/// Calculate Recall@k
///
/// ### Params
///
/// * `true_neighbors` - Slice of true neighbours
/// * `approx_neighbors` - Slice of the approximate neighbours
/// * `k` - Number of selected k
///
/// ### Returns
///
/// The Recall@k
pub fn calculate_recall(
    true_neighbors: &[Vec<usize>],
    approx_neighbors: &[Vec<usize>],
    k: usize,
) -> f64 {
    let mut total_recall = 0.0;

    for (true_nn, approx_nn) in true_neighbors.iter().zip(approx_neighbors.iter()) {
        let true_set: FxHashSet<_> = true_nn.iter().take(k).collect();
        let approx_set: FxHashSet<_> = approx_nn.iter().take(k).collect();

        let matches = approx_set.intersection(&true_set).count();

        total_recall += matches as f64 / k as f64;
    }

    total_recall / true_neighbors.len() as f64
}

/// Calculate mean distance ratio across queries
///
/// For each query, computes the ratio of the sum of approximate distances
/// to the sum of true distances across the top-k neighbours. A ratio of
/// 1.0 indicates perfect results; values above 1.0 indicate how much
/// worse the approximate distances are on average (e.g. 1.05 means 5%
/// worse than optimal).
///
/// This metric is stable across distance metrics and dimensionalities
/// because summing over k neighbours avoids the division-by-near-zero
/// instability that plagues per-pair relative error, particularly with
/// cosine distance. Queries where the true distance sum is negligible
/// (< 1e-12) are excluded.
///
/// ### Params
///
/// * `true_dist` - Slice of true distances to the neighbours (one vec per query)
/// * `approx_dist` - Slice of approximate distances to the neighbours (one vec per query)
/// * `k` - Number of neighbours to consider per query
///
/// ### Returns
///
/// The mean distance ratio (1.0 = perfect, >1.0 = proportionally worse).
/// Returns `NaN` if no queries have a non-negligible true distance sum.
pub fn calculate_mean_distance_ratio<T>(
    true_dist: &[Vec<T>],
    approx_dist: &[Vec<T>],
    k: usize,
) -> f64
where
    T: Float + ToPrimitive,
{
    let mut total_ratio = 0.0;
    let mut count = 0usize;
    for (td, ad) in true_dist.iter().zip(approx_dist.iter()) {
        let n = k.min(td.len()).min(ad.len());
        let sum_true: f64 = td[..n].iter().map(|v| v.to_f64().unwrap()).sum();
        let sum_approx: f64 = ad[..n].iter().map(|v| v.to_f64().unwrap()).sum();
        if sum_true > 1e-12 {
            total_ratio += sum_approx / sum_true;
            count += 1;
        }
    }
    total_ratio / count as f64
}

/// Calculate cluster purity of kNN graph
///
/// Measures what fraction of each point's neighbors belong to the same cluster.
/// High purity (>0.8) means the method preserves cluster structure well.
///
/// ### Params
///
/// * `knn_graph` - Neighbor indices for each point
/// * `cluster_labels` - Ground truth cluster assignment for each point
///
/// ### Returns
///
/// Average fraction of same-cluster neighbors
pub fn calculate_cluster_purity(knn_graph: &[Vec<usize>], cluster_labels: &[usize]) -> f64 {
    let mut total_purity = 0.0;

    for (i, neighbors) in knn_graph.iter().enumerate() {
        let my_cluster = cluster_labels[i];
        let same_cluster = neighbors
            .iter()
            .filter(|&&idx| cluster_labels[idx] == my_cluster)
            .count();
        total_purity += same_cluster as f64 / neighbors.len() as f64;
    }

    total_purity / knn_graph.len() as f64
}

////////////
// Prints //
////////////

fn format_with_underscores(value: f64) -> String {
    let formatted = format!("{:.2}", value);
    let parts: Vec<&str> = formatted.split('.').collect();
    let int_part = parts[0].parse::<i64>().unwrap().separate_with_underscores();
    format!("{}.{}", int_part, parts[1])
}

/// Helper to print results to console
///
/// This version also returns the size of the index
///
/// ### Params
///
/// * `config` - Benchmark configuration
/// * `results` - Benchmark results to print
pub fn print_results_size(config: &str, results: &[BenchmarkResultSize]) {
    println!("\n{:=>131}", "");
    println!("Benchmark: {}", config);
    println!("{:=>131}", "");
    println!(
        "{:<50} {:>12} {:>12} {:>12} {:>12} {:>15} {:>12}",
        "Method",
        "Build (ms)",
        "Query (ms)",
        "Total (ms)",
        "Recall@k",
        "Mean dist ratio",
        "Size (MB)"
    );
    println!("{:->131}", "");
    for result in results {
        println!(
            "{:<50} {:>12} {:>12} {:>12} {:>12.4} {:>15.4} {:>12.2}",
            result.method,
            format_with_underscores(result.build_time_ms),
            format_with_underscores(result.query_time_ms),
            format_with_underscores(result.total_time_ms),
            result.recall_at_k,
            result.mean_dist_rat,
            result.index_size_mb
        );
    }
    println!("{:->131}\n", "");
}

/// Helper to print results to console
///
/// This version prints the cluster purity measure
///
/// ### Params
///
/// * `config` - Benchmark configuration
/// * `results` - Benchmark results to print
pub fn print_results_purity(config: &str, results: &[BenchmarkResultPurity]) {
    println!("\n{:=>128}", "");
    println!("Benchmark: {}", config);
    println!("{:=>128}", "");
    println!(
        "{:<50} {:>12} {:>12} {:>12} {:>12} {:>12} {:>12}",
        "Method", "Build (ms)", "Query (ms)", "Total (ms)", "Recall@k", "Purity", "Size (MB)"
    );
    println!("{:->128}", "");
    for result in results {
        println!(
            "{:<50} {:>12} {:>12} {:>12} {:>12.4} {:>12.4} {:>12.2}",
            result.method,
            format_with_underscores(result.build_time_ms),
            format_with_underscores(result.query_time_ms),
            format_with_underscores(result.total_time_ms),
            result.recall_at_k,
            result.cluster_purity,
            result.index_size_mb
        );
    }
    println!("{:->128}\n", "");
}
