use polars_core::prelude::*;
use std::collections::HashMap;

use crate::error::Result;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Sex {
    Unknown,
    Male,
    Female,
}

impl Sex {
    pub fn from_plink(code: i8) -> Self {
        match code {
            1 => Self::Male,
            2 => Self::Female,
            _ => Self::Unknown,
        }
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Unknown => "unknown",
            Self::Male => "male",
            Self::Female => "female",
        }
    }
}

#[derive(Debug, Clone)]
pub struct SampleInfo {
    pub df: DataFrame,
    index: HashMap<String, usize>,
}

impl SampleInfo {
    pub fn n_samples(&self) -> usize {
        self.df.height()
    }

    pub fn ids(&self) -> Vec<&str> {
        self.df
            .column("sample_id")
            .unwrap()
            .str()
            .unwrap()
            .into_no_null_iter()
            .collect()
    }

    pub fn index_of(&self, id: &str) -> Option<usize> {
        self.index.get(id).copied()
    }

    pub fn subset(&self, ids: &[&str]) -> Result<Self> {
        let mask: Vec<bool> = self
            .ids()
            .iter()
            .map(|sid| ids.contains(sid))
            .collect();
        let mask_col = Column::new("mask".into(), &mask);
        let ca = mask_col.bool().unwrap().clone();
        let df = self.df.filter(&ca)?;
        Ok(Self::from_dataframe(df))
    }

    pub fn from_dataframe(df: DataFrame) -> Self {
        let index = df
            .column("sample_id")
            .unwrap()
            .str()
            .unwrap()
            .into_no_null_iter()
            .enumerate()
            .map(|(i, s)| (s.to_owned(), i))
            .collect();
        Self { df, index }
    }

    pub fn from_plink_fam(
        fid: Vec<String>,
        iid: Vec<String>,
        father: Vec<String>,
        mother: Vec<String>,
        sex: Vec<Sex>,
        pheno: Vec<f64>,
    ) -> Result<Self> {
        let sex_strs: Vec<&str> = sex.iter().map(Sex::as_str).collect();
        let df = DataFrame::new(vec![
            Column::new("sample_id".into(), &iid),
            Column::new("family_id".into(), &fid),
            Column::new("paternal_id".into(), &father),
            Column::new("maternal_id".into(), &mother),
            Column::new("sex".into(), &sex_strs),
            Column::new("phenotype".into(), &pheno),
        ])?;
        Ok(Self::from_dataframe(df))
    }

    pub fn from_ids(ids: Vec<String>) -> Result<Self> {
        let n = ids.len();
        let dots: Vec<&str> = vec!["."; n];
        let zeros: Vec<&str> = vec!["0"; n];
        let unknowns: Vec<&str> = vec!["unknown"; n];
        let nans: Vec<f64> = vec![f64::NAN; n];
        let df = DataFrame::new(vec![
            Column::new("sample_id".into(), &ids),
            Column::new("family_id".into(), &dots),
            Column::new("paternal_id".into(), &zeros),
            Column::new("maternal_id".into(), &zeros),
            Column::new("sex".into(), &unknowns),
            Column::new("phenotype".into(), &nans),
        ])?;
        Ok(Self::from_dataframe(df))
    }
}
