//! Sample datasets for `CellAtlas` demo.
//!
//! Provides generated single-cell datasets for demonstration purposes.

use crate::CellAtlasError;
use cellatlas_shared::{
    Cell, CellCyclePhase, CellState, QualityMetrics, SparseExpression, SpatialCoords,
};

/// Load a sample dataset by ID.
pub fn load_sample_dataset(dataset_id: &str) -> Result<(Vec<Cell>, Vec<String>), CellAtlasError> {
    let gene_names = get_common_gene_names();

    let cells = match dataset_id.to_lowercase().as_str() {
        "pbmc_3k" | "pbmc3k" => generate_pbmc_dataset(2700),
        "mouse_brain_spatial" | "mouse_brain" => generate_spatial_brain_dataset(3500),
        "tumor_microenvironment" | "tumor" => generate_tumor_dataset(8000),
        "developing_heart" | "heart" => generate_heart_dataset(5000),
        _ => {
            return Err(CellAtlasError::InvalidDataset(format!(
                "Unknown dataset: {dataset_id}. Available: pbmc_3k, mouse_brain_spatial, tumor_microenvironment, developing_heart"
            )));
        }
    };

    Ok((cells, gene_names))
}

/// Generate demo cells with specified parameters.
#[must_use]
pub fn generate_demo_cells(n_cells: usize, with_spatial: bool) -> Vec<Cell> {
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal, Poisson, Uniform};

    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
    let n_genes = 20000;

    let n_genes_dist = Normal::new(2000.0_f32, 500.0).unwrap();
    let counts_dist = Normal::new(5000.0_f32, 1500.0).unwrap();
    let mito_dist = Normal::new(5.0_f32, 3.0).unwrap();
    let ribo_dist = Normal::new(10.0_f32, 4.0).unwrap();
    let expr_dist = Poisson::new(2.0_f32).unwrap();
    let spatial_dist = Uniform::new(0.0_f32, 1000.0).unwrap();

    let mut cells = Vec::with_capacity(n_cells);

    for i in 0..n_cells {
        // Generate sparse expression
        let n_expressed = (n_genes_dist.sample(&mut rng) as usize).clamp(100, 5000);
        let gene_indices: Vec<usize> = {
            let mut indices: Vec<usize> = (0..n_genes).collect();
            for j in (1..n_genes).rev() {
                let k = Uniform::new(0, j + 1).unwrap().sample(&mut rng);
                indices.swap(j, k);
            }
            indices.truncate(n_expressed);
            indices.sort_unstable();
            indices
        };

        let values: Vec<f32> = (0..n_expressed)
            .map(|_| (expr_dist.sample(&mut rng) + 1.0).max(0.1))
            .collect();

        let expression = SparseExpression {
            gene_indices,
            values: values.clone(),
            num_genes: n_genes,
        };

        let _total_counts: f32 = values.iter().sum();

        let qc_metrics = QualityMetrics {
            n_genes: n_expressed,
            total_counts: counts_dist.sample(&mut rng).max(100.0),
            pct_mito: mito_dist.sample(&mut rng).clamp(0.0, 50.0),
            pct_ribo: ribo_dist.sample(&mut rng).clamp(0.0, 50.0),
            doublet_score: Some(Uniform::new(0.0_f32, 0.3).unwrap().sample(&mut rng)),
        };

        let spatial_coords = if with_spatial {
            Some(SpatialCoords {
                x: spatial_dist.sample(&mut rng),
                y: spatial_dist.sample(&mut rng),
                z: None,
                section_id: Some(0),
            })
        } else {
            None
        };

        // Assign cell cycle phase based on expression patterns
        let cell_cycle_phase = match i % 5 {
            0 => Some(CellCyclePhase::G1),
            1 => Some(CellCyclePhase::S),
            2 => Some(CellCyclePhase::G2),
            3 => Some(CellCyclePhase::M),
            _ => Some(CellCyclePhase::G0),
        };

        cells.push(Cell {
            id: format!("cell_{i}"),
            barcode: Some(format!("ATCG{i:08}")),
            expression,
            cell_type: None,
            state: Some(CellState {
                name: "Normal".to_string(),
                score: 0.8,
                cell_cycle_phase,
            }),
            qc_metrics,
            spatial_coords,
            cluster_id: None,
            embedding: None,
        });
    }

    cells
}

/// Generate PBMC (Peripheral Blood Mononuclear Cells) dataset.
fn generate_pbmc_dataset(n_cells: usize) -> Vec<Cell> {
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal, Uniform};

    let mut rng = rand::rngs::StdRng::seed_from_u64(12345);
    let gene_names = get_common_gene_names();
    let n_genes = gene_names.len();

    // Cell type proportions for PBMC
    let cell_types = [
        ("CD4+ T cell", 0.30, vec!["CD3D", "CD3E", "CD4", "IL7R"]),
        ("CD8+ T cell", 0.15, vec!["CD3D", "CD3E", "CD8A", "CD8B"]),
        ("B cell", 0.10, vec!["CD19", "MS4A1", "CD79A", "CD79B"]),
        ("Monocyte", 0.20, vec!["CD14", "LYZ", "CST3"]),
        ("NK cell", 0.10, vec!["GNLY", "NKG7", "KLRD1"]),
        ("DC", 0.05, vec!["FCER1A", "CD1C"]),
        ("Platelet", 0.05, vec!["PPBP", "PF4"]),
        ("Other", 0.05, vec![]),
    ];

    let gene_map: std::collections::HashMap<&str, usize> = gene_names
        .iter()
        .enumerate()
        .map(|(i, name)| (name.as_str(), i))
        .collect();

    let mut cells = Vec::with_capacity(n_cells);
    let n_genes_dist = Normal::new(2500.0_f32, 600.0).unwrap();
    let expr_dist = Normal::new(2.0_f32, 1.5).unwrap();

    for i in 0..n_cells {
        // Determine cell type based on proportions
        let r: f32 = Uniform::new(0.0_f32, 1.0).unwrap().sample(&mut rng);
        let mut cumsum = 0.0;
        let mut cell_type_idx = 0;
        for (idx, (_, prop, _)) in cell_types.iter().enumerate() {
            cumsum += prop;
            if r <= cumsum {
                cell_type_idx = idx;
                break;
            }
        }
        let (_, _, markers) = &cell_types[cell_type_idx];

        // Generate expression profile
        let n_expressed = (n_genes_dist.sample(&mut rng) as usize).clamp(500, 4000);
        let mut gene_indices: Vec<usize> = Vec::with_capacity(n_expressed);
        let mut values: Vec<f32> = Vec::with_capacity(n_expressed);

        // Add marker genes with high expression
        for marker in markers {
            if let Some(&idx) = gene_map.get(*marker) {
                gene_indices.push(idx);
                values.push((5.0 + expr_dist.sample(&mut rng)).max(1.0));
            }
        }

        // Add random genes
        let remaining = n_expressed.saturating_sub(gene_indices.len());
        let mut random_genes: Vec<usize> =
            (0..n_genes).filter(|i| !gene_indices.contains(i)).collect();
        for j in (1..random_genes.len().min(remaining + 1)).rev() {
            let k = Uniform::new(0, j + 1).unwrap().sample(&mut rng);
            random_genes.swap(j, k);
        }
        random_genes.truncate(remaining);

        for idx in random_genes {
            gene_indices.push(idx);
            values.push(expr_dist.sample(&mut rng).max(0.1));
        }

        // Sort by gene index
        let mut pairs: Vec<(usize, f32)> = gene_indices
            .iter()
            .copied()
            .zip(values.iter().copied())
            .collect();
        pairs.sort_by_key(|(idx, _)| *idx);
        gene_indices = pairs.iter().map(|(idx, _)| *idx).collect();
        values = pairs.iter().map(|(_, val)| *val).collect();

        let total_counts: f32 = values.iter().sum();

        cells.push(Cell {
            id: format!("pbmc_cell_{i}"),
            barcode: Some(format!("PBMC{i:08}")),
            expression: SparseExpression {
                gene_indices,
                values,
                num_genes: n_genes,
            },
            cell_type: None,
            state: None,
            qc_metrics: QualityMetrics {
                n_genes: n_expressed,
                total_counts,
                pct_mito: Normal::new(4.0_f32, 2.0)
                    .unwrap()
                    .sample(&mut rng)
                    .clamp(0.0, 15.0),
                pct_ribo: Normal::new(12.0_f32, 4.0)
                    .unwrap()
                    .sample(&mut rng)
                    .clamp(0.0, 30.0),
                doublet_score: Some(Uniform::new(0.0_f32, 0.2).unwrap().sample(&mut rng)),
            },
            spatial_coords: None,
            cluster_id: None,
            embedding: None,
        });
    }

    cells
}

/// Generate spatial brain dataset.
fn generate_spatial_brain_dataset(n_cells: usize) -> Vec<Cell> {
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal, Uniform};

    let mut rng = rand::rngs::StdRng::seed_from_u64(54321);
    let gene_names = get_common_gene_names();
    let n_genes = gene_names.len();

    // Brain cell types with spatial organization
    let regions = [
        (
            "Cortex",
            (100.0, 100.0),
            (300.0, 300.0),
            vec!["Neuron", "Astrocyte", "Oligodendrocyte"],
        ),
        (
            "Hippocampus",
            (400.0, 100.0),
            (600.0, 300.0),
            vec!["Neuron", "Microglia"],
        ),
        (
            "Striatum",
            (700.0, 100.0),
            (900.0, 300.0),
            vec!["Neuron", "Astrocyte"],
        ),
        (
            "Thalamus",
            (100.0, 400.0),
            (300.0, 600.0),
            vec!["Neuron", "Oligodendrocyte"],
        ),
    ];

    let mut cells = Vec::with_capacity(n_cells);
    let n_genes_dist = Normal::new(3000.0_f32, 700.0).unwrap();
    let expr_dist = Normal::new(2.5_f32, 1.5).unwrap();

    for i in 0..n_cells {
        // Pick a random region
        let region_idx = Uniform::new(0, regions.len()).unwrap().sample(&mut rng);
        let (_, (x_min, y_min), (x_max, y_max), _cell_types) = &regions[region_idx];

        // Generate spatial coordinates within region
        let x = Uniform::new(*x_min, *x_max).unwrap().sample(&mut rng);
        let y = Uniform::new(*y_min, *y_max).unwrap().sample(&mut rng);

        let n_expressed = (n_genes_dist.sample(&mut rng) as usize).clamp(500, 5000);
        let mut gene_indices: Vec<usize> = (0..n_genes).collect();
        for j in (1..n_genes).rev() {
            let k = Uniform::new(0, j + 1).unwrap().sample(&mut rng);
            gene_indices.swap(j, k);
        }
        gene_indices.truncate(n_expressed);
        gene_indices.sort_unstable();

        let values: Vec<f32> = (0..n_expressed)
            .map(|_| expr_dist.sample(&mut rng).max(0.1))
            .collect();

        let total_counts: f32 = values.iter().sum();

        cells.push(Cell {
            id: format!("brain_cell_{i}"),
            barcode: Some(format!("BRAIN{i:08}")),
            expression: SparseExpression {
                gene_indices,
                values,
                num_genes: n_genes,
            },
            cell_type: None,
            state: None,
            qc_metrics: QualityMetrics {
                n_genes: n_expressed,
                total_counts,
                pct_mito: Normal::new(3.0_f32, 1.5)
                    .unwrap()
                    .sample(&mut rng)
                    .clamp(0.0, 10.0),
                pct_ribo: Normal::new(8.0_f32, 3.0)
                    .unwrap()
                    .sample(&mut rng)
                    .clamp(0.0, 20.0),
                doublet_score: Some(Uniform::new(0.0_f32, 0.15).unwrap().sample(&mut rng)),
            },
            spatial_coords: Some(SpatialCoords {
                x,
                y,
                z: None,
                section_id: Some(0),
            }),
            cluster_id: None,
            embedding: None,
        });
    }

    cells
}

/// Generate tumor microenvironment dataset.
fn generate_tumor_dataset(n_cells: usize) -> Vec<Cell> {
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal, Uniform};

    let mut rng = rand::rngs::StdRng::seed_from_u64(99999);
    let gene_names = get_common_gene_names();
    let n_genes = gene_names.len();

    // TME cell types
    let cell_types = vec![
        ("Tumor cell", 0.40),
        ("TAM", 0.15), // Tumor-associated macrophage
        ("T cell", 0.15),
        ("Fibroblast", 0.10),
        ("Endothelial", 0.10),
        ("B cell", 0.05),
        ("DC", 0.05),
    ];

    let mut cells = Vec::with_capacity(n_cells);
    let n_genes_dist = Normal::new(2800.0_f32, 600.0).unwrap();
    let expr_dist = Normal::new(2.0_f32, 1.2).unwrap();

    for i in 0..n_cells {
        let r: f32 = Uniform::new(0.0_f32, 1.0).unwrap().sample(&mut rng);
        let mut cumsum = 0.0;
        for (_, prop) in &cell_types {
            cumsum += prop;
            if r <= cumsum {
                break;
            }
        }

        let n_expressed = (n_genes_dist.sample(&mut rng) as usize).clamp(500, 5000);
        let mut gene_indices: Vec<usize> = (0..n_genes).collect();
        for j in (1..n_genes).rev() {
            let k = Uniform::new(0, j + 1).unwrap().sample(&mut rng);
            gene_indices.swap(j, k);
        }
        gene_indices.truncate(n_expressed);
        gene_indices.sort_unstable();

        let values: Vec<f32> = (0..n_expressed)
            .map(|_| expr_dist.sample(&mut rng).max(0.1))
            .collect();

        let total_counts: f32 = values.iter().sum();

        cells.push(Cell {
            id: format!("tumor_cell_{i}"),
            barcode: Some(format!("TUMOR{i:08}")),
            expression: SparseExpression {
                gene_indices,
                values,
                num_genes: n_genes,
            },
            cell_type: None,
            state: Some(CellState {
                name: "Activated".to_string(),
                score: Uniform::new(0.5_f32, 1.0).unwrap().sample(&mut rng),
                cell_cycle_phase: Some(match i % 5 {
                    0 => CellCyclePhase::G1,
                    1 => CellCyclePhase::S,
                    2 => CellCyclePhase::G2,
                    3 => CellCyclePhase::M,
                    _ => CellCyclePhase::G0,
                }),
            }),
            qc_metrics: QualityMetrics {
                n_genes: n_expressed,
                total_counts,
                pct_mito: Normal::new(6.0_f32, 3.0)
                    .unwrap()
                    .sample(&mut rng)
                    .clamp(0.0, 20.0),
                pct_ribo: Normal::new(10.0_f32, 4.0)
                    .unwrap()
                    .sample(&mut rng)
                    .clamp(0.0, 25.0),
                doublet_score: Some(Uniform::new(0.0_f32, 0.25).unwrap().sample(&mut rng)),
            },
            spatial_coords: None,
            cluster_id: None,
            embedding: None,
        });
    }

    cells
}

/// Generate developing heart dataset.
fn generate_heart_dataset(n_cells: usize) -> Vec<Cell> {
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal, Uniform};

    let mut rng = rand::rngs::StdRng::seed_from_u64(77777);
    let gene_names = get_common_gene_names();
    let n_genes = gene_names.len();

    // Cardiac cell types
    let cell_types = vec![
        ("Cardiomyocyte", 0.35),
        ("Fibroblast", 0.20),
        ("Endothelial", 0.15),
        ("Smooth muscle", 0.10),
        ("Macrophage", 0.08),
        ("Epicardial", 0.07),
        ("Pericyte", 0.05),
    ];

    let mut cells = Vec::with_capacity(n_cells);
    let n_genes_dist = Normal::new(3200.0_f32, 700.0).unwrap();
    let expr_dist = Normal::new(2.2_f32, 1.4).unwrap();

    for i in 0..n_cells {
        let r: f32 = Uniform::new(0.0_f32, 1.0).unwrap().sample(&mut rng);
        let mut cumsum = 0.0;
        for (_, prop) in &cell_types {
            cumsum += prop;
            if r <= cumsum {
                break;
            }
        }

        let n_expressed = (n_genes_dist.sample(&mut rng) as usize).clamp(600, 5500);
        let mut gene_indices: Vec<usize> = (0..n_genes).collect();
        for j in (1..n_genes).rev() {
            let k = Uniform::new(0, j + 1).unwrap().sample(&mut rng);
            gene_indices.swap(j, k);
        }
        gene_indices.truncate(n_expressed);
        gene_indices.sort_unstable();

        let values: Vec<f32> = (0..n_expressed)
            .map(|_| expr_dist.sample(&mut rng).max(0.1))
            .collect();

        let total_counts: f32 = values.iter().sum();

        cells.push(Cell {
            id: format!("heart_cell_{i}"),
            barcode: Some(format!("HEART{i:08}")),
            expression: SparseExpression {
                gene_indices,
                values,
                num_genes: n_genes,
            },
            cell_type: None,
            state: None,
            qc_metrics: QualityMetrics {
                n_genes: n_expressed,
                total_counts,
                pct_mito: Normal::new(5.0_f32, 2.5)
                    .unwrap()
                    .sample(&mut rng)
                    .clamp(0.0, 15.0),
                pct_ribo: Normal::new(9.0_f32, 3.5)
                    .unwrap()
                    .sample(&mut rng)
                    .clamp(0.0, 22.0),
                doublet_score: Some(Uniform::new(0.0_f32, 0.18).unwrap().sample(&mut rng)),
            },
            spatial_coords: None,
            cluster_id: None,
            embedding: None,
        });
    }

    cells
}

/// Get common gene names for single-cell analysis.
#[must_use]
pub fn get_common_gene_names() -> Vec<String> {
    vec![
        // T cell markers
        "CD3D", "CD3E", "CD3G", "CD4", "CD8A", "CD8B", "IL7R", "CCR7", "SELL", "GZMK", "GZMB",
        "PRF1", "IFNG", "TNF", "IL2", // B cell markers
        "CD19", "MS4A1", "CD79A", "CD79B", "CD27", "CD38", "IGHM", "IGHD",
        // Monocyte/Macrophage markers
        "CD14", "LYZ", "CST3", "FCGR3A", "CD68", "CD163", "MARCO", "MRC1",
        // NK cell markers
        "GNLY", "NKG7", "KLRD1", "KLRB1", "NCAM1", // Dendritic cell markers
        "FCER1A", "CD1C", "CLEC10A", "CD83", "CD86", // Platelet markers
        "PPBP", "PF4", "GP9", "ITGA2B", // Erythrocyte markers
        "HBA1", "HBA2", "HBB", "GYPA", // Housekeeping genes
        "ACTB", "GAPDH", "B2M", "MALAT1", "RPL10", "RPS18", // Mitochondrial genes
        "MT-CO1", "MT-CO2", "MT-ND1", "MT-ATP6", // Stress response
        "JUN", "FOS", "EGR1", "HSP90AA1", // Cell cycle
        "MKI67", "TOP2A", "PCNA", "CDK1", "CCNB1", // Neuron markers
        "SNAP25", "SYT1", "RBFOX3", "MAP2", // Astrocyte markers
        "GFAP", "AQP4", "S100B", // Oligodendrocyte markers
        "MBP", "MOG", "OLIG1", "OLIG2", // Microglia markers
        "AIF1", "CX3CR1", "P2RY12", // Cardiomyocyte markers
        "TNNT2", "MYH7", "ACTC1", "RYR2", // Fibroblast markers
        "COL1A1", "COL1A2", "DCN", "LUM", // Endothelial markers
        "PECAM1", "VWF", "CDH5", "CLDN5",
    ]
    .iter()
    .map(std::string::ToString::to_string)
    .collect()
}

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

    #[test]
    fn test_load_pbmc_dataset() {
        let result = load_sample_dataset("pbmc_3k");
        assert!(result.is_ok());

        let (cells, genes) = result.unwrap();
        assert_eq!(cells.len(), 2700);
        assert!(!genes.is_empty());
    }

    #[test]
    fn test_load_spatial_dataset() {
        let result = load_sample_dataset("mouse_brain_spatial");
        assert!(result.is_ok());

        let (cells, _) = result.unwrap();
        assert_eq!(cells.len(), 3500);
        assert!(cells.iter().all(|c| c.spatial_coords.is_some()));
    }

    #[test]
    fn test_load_tumor_dataset() {
        let result = load_sample_dataset("tumor_microenvironment");
        assert!(result.is_ok());

        let (cells, _) = result.unwrap();
        assert_eq!(cells.len(), 8000);
    }

    #[test]
    fn test_load_heart_dataset() {
        let result = load_sample_dataset("developing_heart");
        assert!(result.is_ok());

        let (cells, _) = result.unwrap();
        assert_eq!(cells.len(), 5000);
    }

    #[test]
    fn test_unknown_dataset() {
        let result = load_sample_dataset("unknown_dataset");
        assert!(result.is_err());
    }

    #[test]
    fn test_generate_demo_cells() {
        let cells = generate_demo_cells(100, false);
        assert_eq!(cells.len(), 100);
        assert!(cells.iter().all(|c| c.spatial_coords.is_none()));

        let spatial_cells = generate_demo_cells(50, true);
        assert_eq!(spatial_cells.len(), 50);
        assert!(spatial_cells.iter().all(|c| c.spatial_coords.is_some()));
    }

    #[test]
    fn test_gene_names() {
        let genes = get_common_gene_names();
        assert!(genes.len() > 50);
        assert!(genes.contains(&"CD3D".to_string()));
        assert!(genes.contains(&"GAPDH".to_string()));
    }

    #[test]
    fn test_cell_quality_metrics() {
        let cells = generate_demo_cells(10, false);
        for cell in &cells {
            assert!(cell.qc_metrics.n_genes > 0);
            assert!(cell.qc_metrics.total_counts > 0.0);
            assert!(cell.qc_metrics.pct_mito >= 0.0 && cell.qc_metrics.pct_mito <= 50.0);
        }
    }
}