use ndarray::{Array2, ArrayView2, Axis};
use polars_core::prelude::*;
use rayon::prelude::*;

use crate::core::{Ploidy, SampleInfo, VariantInfo};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GenotypeValue {
    HomRef,
    Het,
    HomAlt,
    Missing,
}

impl GenotypeValue {
    pub fn to_dosage(self) -> f32 {
        match self {
            Self::HomRef => 0.0,
            Self::Het => 1.0,
            Self::HomAlt => 2.0,
            Self::Missing => f32::NAN,
        }
    }

    pub fn from_plink_code(code: u8) -> Self {
        match code {
            0b00 => Self::HomRef,
            0b01 => Self::Missing,
            0b10 => Self::Het,
            0b11 => Self::HomAlt,
            _ => unreachable!(),
        }
    }

    pub fn from_vcf_gt(gt: &str) -> Self {
        match gt {
            "0/0" | "0|0" => Self::HomRef,
            "0/1" | "0|1" | "1/0" | "1|0" => Self::Het,
            "1/1" | "1|1" => Self::HomAlt,
            _ => Self::Missing,
        }
    }
}

#[derive(Debug, Clone)]
pub struct DosageMatrix {
    data: Array2<f32>,
}

impl DosageMatrix {
    pub fn new(data: Array2<f32>) -> Self {
        Self { data }
    }

    pub fn zeros(n_samples: usize, n_variants: usize) -> Self {
        Self {
            data: Array2::zeros((n_samples, n_variants)),
        }
    }

    pub fn n_samples(&self) -> usize {
        self.data.nrows()
    }

    pub fn n_variants(&self) -> usize {
        self.data.ncols()
    }

    pub fn shape(&self) -> (usize, usize) {
        (self.n_samples(), self.n_variants())
    }

    pub fn as_array(&self) -> ArrayView2<'_, f32> {
        self.data.view()
    }

    pub fn as_array_mut(&mut self) -> &mut Array2<f32> {
        &mut self.data
    }

    pub fn into_array(self) -> Array2<f32> {
        self.data
    }

    #[inline]
    pub fn get(&self, sample: usize, variant: usize) -> f32 {
        self.data[[sample, variant]]
    }

    #[inline]
    pub fn set(&mut self, sample: usize, variant: usize, val: f32) {
        self.data[[sample, variant]] = val;
    }

    pub fn allele_frequencies(&self) -> Vec<f32> {
        (0..self.n_variants())
            .into_par_iter()
            .map(|j: usize| -> f32 {
                let col = self.data.column(j);
                let (sum, count) = col.iter().fold((0.0f64, 0u32), |(s, c), &v| {
                    if v.is_nan() { (s, c) } else { (s + v as f64, c + 1) }
                });
                if count == 0 { f32::NAN } else { (sum / (2.0 * count as f64)) as f32 }
            })
            .collect()
    }

    pub fn missing_rates(&self) -> Vec<f32> {
        let n = self.n_samples() as f32;
        (0..self.n_variants())
            .into_par_iter()
            .map(|j: usize| -> f32 {
                let col = self.data.column(j);
                let n_miss = col.iter().filter(|v| v.is_nan()).count();
                n_miss as f32 / n
            })
            .collect()
    }

    pub fn mean_impute(&mut self) {
        let n_variants = self.n_variants();
        let n_samples = self.n_samples();

        // Compute means in parallel
        let means: Vec<f32> = (0..n_variants)
            .into_par_iter()
            .map(|j: usize| -> f32 {
                let col = self.data.column(j);
                let (sum, count) = col.iter().fold((0.0f64, 0u32), |(s, c), &v| {
                    if v.is_nan() { (s, c) } else { (s + v as f64, c + 1) }
                });
                if count == 0 { 0.0 } else { (sum / count as f64) as f32 }
            })
            .collect();

        // Impute — simple loops (writes are cheap, parallelism overhead not worth it)
        for j in 0..n_variants {
            let mean = means[j];
            for i in 0..n_samples {
                if self.data[[i, j]].is_nan() {
                    self.data[[i, j]] = mean;
                }
            }
        }
    }

    pub fn select_variants(&self, indices: &[usize]) -> Self {
        let data = self.data.select(Axis(1), indices);
        Self { data }
    }

    pub fn select_samples(&self, indices: &[usize]) -> Self {
        let data = self.data.select(Axis(0), indices);
        Self { data }
    }
}

/// The final, unified output of every bunbun parser.
#[derive(Debug, Clone)]
pub struct GenotypeData {
    pub samples: SampleInfo,
    pub variants: VariantInfo,
    pub dosages: DosageMatrix,
    pub ploidy: Ploidy,
    pub source_format: &'static str,
}

impl GenotypeData {
    pub fn n_samples(&self) -> usize {
        self.samples.n_samples()
    }

    pub fn n_variants(&self) -> usize {
        self.variants.n_variants()
    }

    pub fn shape(&self) -> (usize, usize) {
        self.dosages.shape()
    }

    pub fn subset_samples(&self, ids: &[&str]) -> crate::error::Result<Self> {
        let indices: Vec<usize> = ids
            .iter()
            .filter_map(|id| self.samples.index_of(id))
            .collect();
        let samples = self.samples.subset(ids)?;
        let dosages = self.dosages.select_samples(&indices);
        Ok(Self {
            samples,
            variants: self.variants.clone(),
            dosages,
            ploidy: self.ploidy,
            source_format: self.source_format,
        })
    }

    pub fn region(
        &self,
        chrom: &str,
        start: u32,
        end: u32,
    ) -> crate::error::Result<Self> {
        let indices = self.variants.region(chrom, start, end)?;

        // Build boolean mask and filter using Column (works across polars-core versions)
        let mask: Vec<bool> = (0..self.n_variants())
            .map(|i| indices.contains(&i))
            .collect();
        let mask_col = Column::new("m".into(), &mask);
        let ca = mask_col.bool().unwrap().clone();
        let vdf = self.variants.df.filter(&ca)?;

        let alleles: Vec<_> = indices
            .iter()
            .map(|&i| {
                let (r, a) = self.variants.alleles_at(i);
                (r.clone(), a.to_vec())
            })
            .collect();

        let new_variants = VariantInfo::from_parts(
            vdf.column("chrom")?.str()?.into_no_null_iter().map(|s| s.to_owned()).collect(),
            vdf.column("pos")?.u32()?.into_no_null_iter().collect(),
            vdf.column("id")?.str()?.into_no_null_iter().map(|s| s.to_owned()).collect(),
            alleles.iter().map(|(r, _)| r.clone()).collect(),
            alleles.iter().map(|(_, a)| a.clone()).collect(),
            vdf.column("cm")?.f64()?.into_no_null_iter().collect(),
            vdf.column("qual")?.f64()?.into_no_null_iter().collect(),
            vdf.column("filter")?.str()?.into_no_null_iter().map(|s| s.to_owned()).collect(),
        )?;

        let dosages = self.dosages.select_variants(&indices);

        Ok(Self {
            samples: self.samples.clone(),
            variants: new_variants,
            dosages,
            ploidy: self.ploidy,
            source_format: self.source_format,
        })
    }
}
