use polars_core::prelude::*;

use crate::core::Allele;
use crate::error::Result;

/// Variant metadata for a set of genetic variants.
#[derive(Debug, Clone)]
pub struct VariantInfo {
    pub df: DataFrame,
    ref_alleles: Vec<Allele>,
    alt_alleles: Vec<Vec<Allele>>,
}

impl VariantInfo {
    pub fn n_variants(&self) -> usize {
        self.df.height()
    }

    pub fn alleles_at(&self, idx: usize) -> (&Allele, &[Allele]) {
        (&self.ref_alleles[idx], &self.alt_alleles[idx])
    }

    /// Return indices of variants within the genomic region [start, end].
    pub fn region(&self, chrom: &str, start: u32, end: u32) -> Result<Vec<usize>> {
        let chroms = self.df.column("chrom")?.str()?;
        let positions = self.df.column("pos")?.u32()?;
        let mut indices = Vec::new();
        for i in 0..self.n_variants() {
            if let (Some(c), Some(p)) = (chroms.get(i), positions.get(i)) {
                if c == chrom && p >= start && p <= end {
                    indices.push(i);
                }
            }
        }
        Ok(indices)
    }

    /// Construct from PLINK .bim columns.
    pub fn from_bim(
        chrom: Vec<String>,
        id: Vec<String>,
        cm: Vec<f64>,
        pos: Vec<u32>,
        a1: Vec<String>,
        a2: Vec<String>,
    ) -> Result<Self> {
        let n = chrom.len();
        let ref_alleles: Vec<Allele> = a1.iter().map(|s| Allele::new(s)).collect();
        let alt_alleles: Vec<Vec<Allele>> = a2.iter().map(|s| vec![Allele::new(s)]).collect();

        let df = DataFrame::new(vec![
            Series::new("chrom".into(), &chrom).into(),
            Series::new("pos".into(), &pos).into(),
            Series::new("id".into(), &id).into(),
            Series::new("ref".into(), &a1).into(),
            Series::new("alt".into(), &a2).into(),
            Series::new("cm".into(), &cm).into(),
            Series::new("qual".into(), &vec![f64::NAN; n]).into(),
            Series::new("filter".into(), &vec!["PASS".to_owned(); n]).into(),
        ])?;

        Ok(Self { df, ref_alleles, alt_alleles })
    }

    /// Construct from already-parsed columns (VCF, BGEN, PLINK2).
    pub fn from_parts(
        chrom: Vec<String>,
        pos: Vec<u32>,
        id: Vec<String>,
        ref_alleles: Vec<Allele>,
        alt_alleles: Vec<Vec<Allele>>,
        cm: Vec<f64>,
        qual: Vec<f64>,
        filter: Vec<String>,
    ) -> Result<Self> {
        let ref_strs: Vec<String> = ref_alleles.iter().map(|a| a.as_str().to_owned()).collect();
        let alt_strs: Vec<String> = alt_alleles
            .iter()
            .map(|alts| {
                alts.iter()
                    .map(|a| a.as_str().to_owned())
                    .collect::<Vec<_>>()
                    .join(",")
            })
            .collect();

        let df = DataFrame::new(vec![
            Series::new("chrom".into(), &chrom).into(),
            Series::new("pos".into(), &pos).into(),
            Series::new("id".into(), &id).into(),
            Series::new("ref".into(), &ref_strs).into(),
            Series::new("alt".into(), &alt_strs).into(),
            Series::new("cm".into(), &cm).into(),
            Series::new("qual".into(), &qual).into(),
            Series::new("filter".into(), &filter).into(),
        ])?;

        Ok(Self { df, ref_alleles, alt_alleles })
    }
}
