use bunbun::core::*;
use bunbun::io::FormatDetector;
use ndarray::Array2;
use std::path::Path;

#[test]
fn allele_inline_vs_heap() {
    let short = Allele::new("ATG");
    assert_eq!(short.as_str(), "ATG");
    assert!(!short.is_missing());

    let long = Allele::new("ATCGATCGATCG");
    assert_eq!(long.as_str(), "ATCGATCGATCG");

    let missing = Allele::new(".");
    assert!(missing.is_missing());
}

#[test]
fn dosage_matrix_stats() {
    // 3 samples × 4 variants
    let data = Array2::from_shape_vec(
        (3, 4),
        vec![
            0.0, 1.0, 2.0, f32::NAN,
            0.0, 0.0, 1.0, 1.0,
            2.0, 1.0, 0.0, 2.0,
        ],
    ).unwrap();

    let mat = DosageMatrix::new(data);
    assert_eq!(mat.shape(), (3, 4));

    let freqs = mat.allele_frequencies();
    assert_eq!(freqs.len(), 4);
    // Variant 0: (0+0+2)/(2*3) = 2/6 ≈ 0.333
    assert!((freqs[0] - 0.333).abs() < 0.01);
    // Variant 1: (1+0+1)/(2*3) = 2/6 ≈ 0.333
    assert!((freqs[1] - 0.333).abs() < 0.01);

    let miss = mat.missing_rates();
    assert_eq!(miss.len(), 4);
    assert!((miss[3] - 1.0 / 3.0).abs() < 0.01); // 1 of 3 missing
    assert!((miss[0] - 0.0).abs() < 0.001);       // 0 missing
}

#[test]
fn dosage_matrix_mean_impute() {
    let data = Array2::from_shape_vec(
        (3, 2),
        vec![
            0.0, f32::NAN,
            1.0, 2.0,
            2.0, 1.0,
        ],
    ).unwrap();

    let mut mat = DosageMatrix::new(data);
    mat.mean_impute();

    let arr = mat.as_array();
    // Variant 0: mean of (0, 1, 2) = 1.0 → no NaN to impute
    assert_eq!(arr[[0, 0]], 0.0);
    // Variant 1: mean of (2, 1) = 1.5 → NaN at [0,1] → 1.5
    assert!((arr[[0, 1]] - 1.5).abs() < 0.001);
}

#[test]
fn dosage_matrix_subsetting() {
    let data = Array2::from_shape_vec(
        (4, 3),
        vec![
            0.0, 1.0, 2.0,
            1.0, 1.0, 0.0,
            2.0, 0.0, 1.0,
            0.0, 2.0, 0.0,
        ],
    ).unwrap();

    let mat = DosageMatrix::new(data);

    let sub_v = mat.select_variants(&[0, 2]);
    assert_eq!(sub_v.shape(), (4, 2));
    assert_eq!(sub_v.get(0, 0), 0.0);
    assert_eq!(sub_v.get(0, 1), 2.0);

    let sub_s = mat.select_samples(&[1, 3]);
    assert_eq!(sub_s.shape(), (2, 3));
    assert_eq!(sub_s.get(0, 0), 1.0); // was row 1
    assert_eq!(sub_s.get(1, 1), 2.0); // was row 3, col 1
}

#[test]
fn genotype_value_plink_encoding() {
    assert_eq!(GenotypeValue::from_plink_code(0b00).to_dosage(), 0.0);
    assert_eq!(GenotypeValue::from_plink_code(0b10).to_dosage(), 1.0);
    assert_eq!(GenotypeValue::from_plink_code(0b11).to_dosage(), 2.0);
    assert!(GenotypeValue::from_plink_code(0b01).to_dosage().is_nan());
}

#[test]
fn genotype_value_vcf_gt() {
    assert_eq!(GenotypeValue::from_vcf_gt("0/0").to_dosage(), 0.0);
    assert_eq!(GenotypeValue::from_vcf_gt("0|1").to_dosage(), 1.0);
    assert_eq!(GenotypeValue::from_vcf_gt("1/0").to_dosage(), 1.0);
    assert_eq!(GenotypeValue::from_vcf_gt("1|1").to_dosage(), 2.0);
    assert!(GenotypeValue::from_vcf_gt("./.").to_dosage().is_nan());
}

#[test]
fn format_detection() {
    assert_eq!(
        FormatDetector::detect(Path::new("data.bed")).unwrap(),
        bunbun::io::Format::Plink
    );
    assert_eq!(
        FormatDetector::detect(Path::new("data.vcf")).unwrap(),
        bunbun::io::Format::Vcf
    );
    assert_eq!(
        FormatDetector::detect(Path::new("data.vcf.gz")).unwrap(),
        bunbun::io::Format::Vcf
    );
    assert_eq!(
        FormatDetector::detect(Path::new("data.pgen")).unwrap(),
        bunbun::io::Format::Plink2
    );
    assert_eq!(
        FormatDetector::detect(Path::new("data.bgen")).unwrap(),
        bunbun::io::Format::Bgen
    );
}

#[test]
fn plink_path_resolution() {
    let (bed, bim, fam) = FormatDetector::resolve_plink(Path::new("/data/cohort.bed"));
    assert_eq!(bed, Path::new("/data/cohort.bed"));
    assert_eq!(bim, Path::new("/data/cohort.bim"));
    assert_eq!(fam, Path::new("/data/cohort.fam"));
}

#[test]
fn sample_info_from_ids() {
    let ids = vec!["S1".into(), "S2".into(), "S3".into()];
    let info = SampleInfo::from_ids(ids).unwrap();
    assert_eq!(info.n_samples(), 3);
    assert_eq!(info.ids(), vec!["S1", "S2", "S3"]);
    assert_eq!(info.index_of("S2"), Some(1));
    assert_eq!(info.index_of("S99"), None);
}

#[test]
fn ploidy_values() {
    assert_eq!(Ploidy::Haploid.n(), 1);
    assert_eq!(Ploidy::Diploid.n(), 2);
    assert_eq!(Ploidy::Polyploid(6).n(), 6);
    assert_eq!(Ploidy::default(), Ploidy::Diploid);
}
