//! Cell type annotation algorithms.
//!
//! Implements marker-based annotation, transfer learning, and
//! foundation model-based cell type prediction.

use crate::CellAtlasError;
use cellatlas_shared::{
    AnalysisConfig, AnnotationMethod, Cell, CellType, CellTypeAnnotation, ClusterInfo,
    MarkerGeneSet, get_common_markers,
};

/// Annotate cells with cell types.
pub fn annotate_cells(
    cells: &[Cell],
    clusters: &[ClusterInfo],
    config: &AnalysisConfig,
) -> Result<(Vec<Cell>, Vec<CellTypeAnnotation>), CellAtlasError> {
    match config.annotation.method {
        AnnotationMethod::MarkerBased => {
            let markers = config
                .annotation
                .marker_genes
                .clone()
                .unwrap_or_else(get_common_markers);
            marker_based_annotation(cells, clusters, &markers)
        }
        AnnotationMethod::TransferLearning => {
            transfer_learning_annotation(cells, clusters, &config.annotation.reference)
        }
        AnnotationMethod::FoundationModel => foundation_model_annotation(cells, clusters),
    }
}

/// Marker-based cell type annotation.
///
/// Assigns cell types based on expression of known marker genes.
fn marker_based_annotation(
    cells: &[Cell],
    _clusters: &[ClusterInfo],
    marker_sets: &[MarkerGeneSet],
) -> Result<(Vec<Cell>, Vec<CellTypeAnnotation>), CellAtlasError> {
    let mut annotated_cells = cells.to_vec();
    let mut cell_type_counts: std::collections::HashMap<String, CellTypeStats> =
        std::collections::HashMap::new();

    // Build gene name to index map (simulated)
    let gene_map = build_gene_map();

    // Annotate each cell
    for cell in &mut annotated_cells {
        let best_annotation = find_best_cell_type(cell, marker_sets, &gene_map);
        cell.cell_type = best_annotation.clone();

        if let Some(ct) = best_annotation {
            let stats = cell_type_counts
                .entry(ct.name.clone())
                .or_insert_with(|| CellTypeStats {
                    count: 0,
                    total_confidence: 0.0,
                    clusters: std::collections::HashSet::new(),
                });
            stats.count += 1;
            stats.total_confidence += ct.confidence;
            if let Some(cluster_id) = cell.cluster_id {
                stats.clusters.insert(cluster_id);
            }
        }
    }

    // Build cell type annotations summary
    let total_cells = annotated_cells.len();
    let cell_types: Vec<CellTypeAnnotation> = cell_type_counts
        .into_iter()
        .map(|(name, stats)| CellTypeAnnotation {
            name,
            n_cells: stats.count,
            percentage: (stats.count as f32 / total_cells as f32) * 100.0,
            avg_confidence: stats.total_confidence / stats.count as f32,
            clusters: stats.clusters.into_iter().collect(),
        })
        .collect();

    Ok((annotated_cells, cell_types))
}

/// Transfer learning annotation using a reference dataset.
fn transfer_learning_annotation(
    cells: &[Cell],
    _clusters: &[ClusterInfo],
    reference: &Option<String>,
) -> Result<(Vec<Cell>, Vec<CellTypeAnnotation>), CellAtlasError> {
    // For demo, simulate transfer learning
    let reference_name = reference.as_deref().unwrap_or("Human Cell Atlas");
    tracing::info!("Using reference dataset: {}", reference_name);

    // Create reference cell type signatures (simulated)
    let reference_signatures = get_reference_signatures();

    let mut annotated_cells = cells.to_vec();
    let mut cell_type_counts: std::collections::HashMap<String, CellTypeStats> =
        std::collections::HashMap::new();

    for cell in &mut annotated_cells {
        // Match cell to reference signatures
        let best_match = match_to_reference(cell, &reference_signatures);
        cell.cell_type = Some(best_match.clone());

        let stats = cell_type_counts
            .entry(best_match.name.clone())
            .or_insert_with(|| CellTypeStats {
                count: 0,
                total_confidence: 0.0,
                clusters: std::collections::HashSet::new(),
            });
        stats.count += 1;
        stats.total_confidence += best_match.confidence;
        if let Some(cluster_id) = cell.cluster_id {
            stats.clusters.insert(cluster_id);
        }
    }

    let total_cells = annotated_cells.len();
    let cell_types: Vec<CellTypeAnnotation> = cell_type_counts
        .into_iter()
        .map(|(name, stats)| CellTypeAnnotation {
            name,
            n_cells: stats.count,
            percentage: (stats.count as f32 / total_cells as f32) * 100.0,
            avg_confidence: stats.total_confidence / stats.count as f32,
            clusters: stats.clusters.into_iter().collect(),
        })
        .collect();

    Ok((annotated_cells, cell_types))
}

/// Foundation model-based annotation using `CellAtlas` transformer.
fn foundation_model_annotation(
    cells: &[Cell],
    clusters: &[ClusterInfo],
) -> Result<(Vec<Cell>, Vec<CellTypeAnnotation>), CellAtlasError> {
    use crate::cell_transformer::{CellTransformer, CellTransformerConfig};

    tracing::info!("Using CellAtlas foundation model for annotation");

    // Initialize transformer
    let config = CellTransformerConfig {
        num_genes: 2000,
        hidden_dim: 128,
        num_heads: 4,
        num_layers: 4,
        dropout: 0.1,
        max_seq_len: 512,
    };
    let transformer = CellTransformer::new(config);

    // Encode cells
    let gene_indices: Vec<Vec<usize>> = cells
        .iter()
        .map(|c| c.expression.gene_indices.clone())
        .collect();
    let values: Vec<Vec<f32>> = cells.iter().map(|c| c.expression.values.clone()).collect();

    let embeddings = transformer.encode(&gene_indices, &values);

    // Classify based on embeddings (simulated classifier)
    let cell_type_predictions = classify_embeddings(&embeddings, clusters);

    let mut annotated_cells = cells.to_vec();
    let mut cell_type_counts: std::collections::HashMap<String, CellTypeStats> =
        std::collections::HashMap::new();

    for (i, cell) in annotated_cells.iter_mut().enumerate() {
        if i < cell_type_predictions.len() {
            let prediction = &cell_type_predictions[i];
            cell.cell_type = Some(prediction.clone());

            let stats = cell_type_counts
                .entry(prediction.name.clone())
                .or_insert_with(|| CellTypeStats {
                    count: 0,
                    total_confidence: 0.0,
                    clusters: std::collections::HashSet::new(),
                });
            stats.count += 1;
            stats.total_confidence += prediction.confidence;
            if let Some(cluster_id) = cell.cluster_id {
                stats.clusters.insert(cluster_id);
            }
        }
    }

    let total_cells = annotated_cells.len();
    let cell_types: Vec<CellTypeAnnotation> = cell_type_counts
        .into_iter()
        .map(|(name, stats)| CellTypeAnnotation {
            name,
            n_cells: stats.count,
            percentage: (stats.count as f32 / total_cells as f32) * 100.0,
            avg_confidence: stats.total_confidence / stats.count as f32,
            clusters: stats.clusters.into_iter().collect(),
        })
        .collect();

    Ok((annotated_cells, cell_types))
}

/// Helper struct for tracking cell type statistics.
#[derive(Debug)]
struct CellTypeStats {
    count: usize,
    total_confidence: f32,
    clusters: std::collections::HashSet<usize>,
}

/// Build gene name to index map.
fn build_gene_map() -> std::collections::HashMap<String, usize> {
    let common_genes = vec![
        "CD3D", "CD3E", "CD3G", "CD4", "CD8A", "CD8B", "CD19", "MS4A1", "CD79A", "CD79B", "CD14",
        "LYZ", "CST3", "FCGR3A", "GNLY", "NKG7", "KLRD1", "PRF1", "FCER1A", "CD1C", "CLEC10A",
        "IL7R", "CCR7", "S100A4", "PPBP", "PF4", "GP9", "HBA1", "HBA2", "HBB",
    ];

    common_genes
        .iter()
        .enumerate()
        .map(|(i, &gene)| (gene.to_string(), i))
        .collect()
}

/// Find best cell type for a cell based on marker expression.
fn find_best_cell_type(
    cell: &Cell,
    marker_sets: &[MarkerGeneSet],
    gene_map: &std::collections::HashMap<String, usize>,
) -> Option<CellType> {
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal};

    let mut rng = rand::rngs::StdRng::seed_from_u64(cell.id.len() as u64);
    let noise = Normal::new(0.0_f32, 0.1).unwrap();

    let mut best_score = 0.0_f32;
    let mut best_type: Option<CellType> = None;

    for marker_set in marker_sets {
        // Count positive marker expression
        let mut positive_count = 0;
        let mut positive_total = 0.0;

        for marker in &marker_set.positive_markers {
            if let Some(&gene_idx) = gene_map.get(marker)
                && gene_idx < cell.expression.gene_indices.len()
            {
                let expr = cell.expression.get(gene_idx);
                if expr > 0.0 {
                    positive_count += 1;
                    positive_total += expr;
                }
            }
        }

        // Check negative markers
        let mut negative_count = 0;
        for marker in &marker_set.negative_markers {
            if let Some(&gene_idx) = gene_map.get(marker) {
                let expr = cell.expression.get(gene_idx);
                if expr > 0.0 {
                    negative_count += 1;
                }
            }
        }

        // Calculate score
        let positive_score = if marker_set.positive_markers.is_empty() {
            0.5
        } else {
            positive_count as f32 / marker_set.positive_markers.len() as f32
        };

        let negative_penalty = if marker_set.negative_markers.is_empty() {
            0.0
        } else {
            negative_count as f32 / marker_set.negative_markers.len() as f32 * 0.5
        };

        let score = (positive_score - negative_penalty + noise.sample(&mut rng)).clamp(0.0, 1.0);

        if score > best_score {
            best_score = score;
            best_type = Some(CellType {
                name: marker_set.cell_type.clone(),
                confidence: score,
                lineage: vec!["Immune".to_string(), marker_set.cell_type.clone()],
                markers: marker_set.positive_markers.clone(),
            });
        }
    }

    best_type
}

/// Get reference cell type signatures for transfer learning.
fn get_reference_signatures() -> Vec<ReferenceSignature> {
    vec![
        ReferenceSignature {
            cell_type: "CD4+ T cell".to_string(),
            marker_weights: vec![
                ("CD3D".to_string(), 1.0),
                ("CD4".to_string(), 1.0),
                ("IL7R".to_string(), 0.8),
            ],
            lineage: vec![
                "Immune".to_string(),
                "T cell".to_string(),
                "CD4+ T cell".to_string(),
            ],
        },
        ReferenceSignature {
            cell_type: "CD8+ T cell".to_string(),
            marker_weights: vec![
                ("CD3D".to_string(), 1.0),
                ("CD8A".to_string(), 1.0),
                ("GZMK".to_string(), 0.7),
            ],
            lineage: vec![
                "Immune".to_string(),
                "T cell".to_string(),
                "CD8+ T cell".to_string(),
            ],
        },
        ReferenceSignature {
            cell_type: "B cell".to_string(),
            marker_weights: vec![
                ("CD19".to_string(), 1.0),
                ("MS4A1".to_string(), 1.0),
                ("CD79A".to_string(), 0.9),
            ],
            lineage: vec!["Immune".to_string(), "B cell".to_string()],
        },
        ReferenceSignature {
            cell_type: "Monocyte".to_string(),
            marker_weights: vec![
                ("CD14".to_string(), 1.0),
                ("LYZ".to_string(), 0.9),
                ("CST3".to_string(), 0.8),
            ],
            lineage: vec![
                "Immune".to_string(),
                "Myeloid".to_string(),
                "Monocyte".to_string(),
            ],
        },
        ReferenceSignature {
            cell_type: "NK cell".to_string(),
            marker_weights: vec![
                ("GNLY".to_string(), 1.0),
                ("NKG7".to_string(), 1.0),
                ("KLRD1".to_string(), 0.8),
            ],
            lineage: vec!["Immune".to_string(), "NK cell".to_string()],
        },
        ReferenceSignature {
            cell_type: "Dendritic cell".to_string(),
            marker_weights: vec![("FCER1A".to_string(), 1.0), ("CD1C".to_string(), 0.9)],
            lineage: vec![
                "Immune".to_string(),
                "Myeloid".to_string(),
                "Dendritic cell".to_string(),
            ],
        },
    ]
}

/// Reference cell type signature.
#[derive(Debug)]
struct ReferenceSignature {
    cell_type: String,
    marker_weights: Vec<(String, f32)>,
    lineage: Vec<String>,
}

/// Match a cell to reference signatures.
fn match_to_reference(cell: &Cell, references: &[ReferenceSignature]) -> CellType {
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal};

    let mut rng = rand::rngs::StdRng::seed_from_u64(cell.id.len() as u64 * 17);
    let noise = Normal::new(0.0_f32, 0.1).unwrap();

    let gene_map = build_gene_map();

    let mut best_score = 0.0_f32;
    let mut best_ref = &references[0];

    for reference in references {
        let mut score = 0.0_f32;
        let mut total_weight = 0.0_f32;

        for (gene, weight) in &reference.marker_weights {
            if let Some(&gene_idx) = gene_map.get(gene) {
                let expr = cell.expression.get(gene_idx);
                if expr > 0.0 {
                    score += weight;
                }
                total_weight += weight;
            }
        }

        let normalized_score = if total_weight > 0.0 {
            score / total_weight
        } else {
            0.0
        };

        let noisy_score = (normalized_score + noise.sample(&mut rng)).clamp(0.0, 1.0);

        if noisy_score > best_score {
            best_score = noisy_score;
            best_ref = reference;
        }
    }

    CellType {
        name: best_ref.cell_type.clone(),
        confidence: best_score,
        lineage: best_ref.lineage.clone(),
        markers: best_ref
            .marker_weights
            .iter()
            .map(|(g, _)| g.clone())
            .collect(),
    }
}

/// Classify cell embeddings using a simple classifier.
fn classify_embeddings(embeddings: &[Vec<f32>], _clusters: &[ClusterInfo]) -> Vec<CellType> {
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal};

    let cell_types = [
        ("CD4+ T cell", vec!["Immune", "T cell", "CD4+ T cell"]),
        ("CD8+ T cell", vec!["Immune", "T cell", "CD8+ T cell"]),
        ("B cell", vec!["Immune", "B cell"]),
        ("Monocyte", vec!["Immune", "Myeloid", "Monocyte"]),
        ("NK cell", vec!["Immune", "NK cell"]),
        ("Dendritic cell", vec!["Immune", "Myeloid", "DC"]),
    ];

    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
    let confidence_noise = Normal::new(0.0_f32, 0.1).unwrap();

    embeddings
        .iter()
        .enumerate()
        .map(|(i, emb)| {
            // Simple classification based on embedding features
            let sum: f32 = emb.iter().sum();
            let type_idx = (i + (sum.abs() as usize)) % cell_types.len();
            let (name, lineage) = &cell_types[type_idx];

            CellType {
                name: name.to_string(),
                confidence: (0.75 + confidence_noise.sample(&mut rng)).clamp(0.5, 0.99),
                lineage: lineage
                    .iter()
                    .map(std::string::ToString::to_string)
                    .collect(),
                markers: vec![],
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use cellatlas_shared::{AnnotationConfig, Embedding2D, QualityMetrics, SparseExpression};

    fn create_test_cell(id: &str, gene_indices: Vec<usize>, values: Vec<f32>) -> Cell {
        Cell {
            id: id.to_string(),
            barcode: None,
            expression: SparseExpression {
                gene_indices,
                values: values.clone(),
                num_genes: 100,
            },
            cell_type: None,
            state: None,
            qc_metrics: QualityMetrics {
                n_genes: values.len(),
                total_counts: values.iter().sum(),
                pct_mito: 3.0,
                pct_ribo: 10.0,
                doublet_score: None,
            },
            spatial_coords: None,
            cluster_id: Some(0),
            embedding: Some(Embedding2D { x: 0.0, y: 0.0 }),
        }
    }

    #[test]
    fn test_marker_based_annotation() {
        let cells = vec![
            create_test_cell("t_cell", vec![0, 1, 2], vec![5.0, 4.0, 3.0]),
            create_test_cell("b_cell", vec![6, 7, 8], vec![6.0, 5.0, 4.0]),
        ];

        let clusters = vec![ClusterInfo {
            id: 0,
            n_cells: 2,
            markers: vec!["CD3D".to_string()],
            centroid: None,
        }];

        let markers = get_common_markers();

        let result = marker_based_annotation(&cells, &clusters, &markers);
        assert!(result.is_ok());

        let (annotated, cell_types) = result.unwrap();
        assert_eq!(annotated.len(), 2);
        assert!(!cell_types.is_empty());
    }

    #[test]
    fn test_transfer_learning_annotation() {
        let cells = vec![
            create_test_cell("cell_1", vec![0, 1], vec![5.0, 4.0]),
            create_test_cell("cell_2", vec![6, 7], vec![6.0, 5.0]),
        ];

        let clusters = vec![ClusterInfo {
            id: 0,
            n_cells: 2,
            markers: vec![],
            centroid: None,
        }];

        let result = transfer_learning_annotation(&cells, &clusters, &None);
        assert!(result.is_ok());

        let (annotated, cell_types) = result.unwrap();
        assert!(annotated.iter().all(|c| c.cell_type.is_some()));
        assert!(!cell_types.is_empty());
    }

    #[test]
    fn test_foundation_model_annotation() {
        let cells = vec![
            create_test_cell("cell_1", vec![0, 1, 2], vec![1.0, 2.0, 1.5]),
            create_test_cell("cell_2", vec![3, 4, 5], vec![2.0, 1.0, 0.5]),
        ];

        let clusters = vec![ClusterInfo {
            id: 0,
            n_cells: 2,
            markers: vec![],
            centroid: None,
        }];

        let result = foundation_model_annotation(&cells, &clusters);
        assert!(result.is_ok());

        let (annotated, cell_types) = result.unwrap();
        assert_eq!(annotated.len(), 2);
    }

    #[test]
    fn test_reference_signatures() {
        let signatures = get_reference_signatures();
        assert!(!signatures.is_empty());
        assert!(signatures.iter().any(|s| s.cell_type.contains("T cell")));
    }

    #[test]
    fn test_build_gene_map() {
        let map = build_gene_map();
        assert!(map.contains_key("CD3D"));
        assert!(map.contains_key("CD19"));
        assert!(map.contains_key("GNLY"));
    }
}