use anyhow::{anyhow, Ok, Result};
use rust_htslib::bam::IndexedReader;
use std::{cmp::Ordering, collections::HashSet};

#[derive(Debug, Clone)]
pub struct HeteroVar {
    pub chr: String,
    pub position: i32,
    pub base: u8,
    pub vaf: f64,
    pub reads: HashSet<String>,
}

impl HeteroVar {
    pub fn new(
        bam: &mut IndexedReader,
        chr: &str,
        position: i32,
        reference: u8,
        alternative: u8,
    ) -> Result<(HeteroVar, HeteroVar)> {
        let rec_base = if let std::result::Result::Ok(rb) =
            pileup::qnames_at_base(bam, chr, position, false)
        {
            rb
        } else {
            return Err(anyhow!("Error while reading BAM file."));
        };
        let depth = rec_base.len() as i32;
        if depth == 0 {
            return Err(anyhow!("No records"));
        }
        let mut a = (reference, Vec::new());
        let mut b = (alternative, Vec::new());
        for (record, base) in rec_base.into_iter() {
            if base == reference {
                a.1.push(record.clone());
            } else if base == alternative {
                b.1.push(record.clone());
            }
        }

        let res: Vec<HeteroVar> = vec![a, b]
            .into_iter()
            .enumerate()
            .filter(|(_, (_, rec))| rec.len() > 0)
            .map(|(_, (base, records))| {
                let mut records_ids: HashSet<String> = HashSet::new();
                records_ids.extend(records.into_iter());
                HeteroVar {
                    chr: chr.to_string(),
                    position,
                    base,
                    vaf: records_ids.len() as f64 / depth as f64,
                    reads: records_ids,
                }
            })
            .collect();

        if res.len() == 2 {
            let a = res.get(0).unwrap().to_owned();
            let b = res.get(1).unwrap().to_owned();
            return Ok((a, b));
        } else {
            return Err(anyhow!("Not enough reads"));
        }
    }
}
impl PartialEq for HeteroVar {
    fn eq(&self, other: &Self) -> bool {
        self.reads == other.reads
            && self.chr == other.chr
            && self.position == other.position
            && self.base == other.base
    }
}

#[derive(Debug, Clone)]
pub struct Phase {
    pub data: Vec<HeteroVar>,
    pub reads: HashSet<String>,
}

impl Phase {
    pub fn new(hetero_var: &HeteroVar) -> Self {
        Self {
            data: vec![hetero_var.clone()],
            reads: hetero_var.reads.clone(),
        }
    }
    pub fn try_merge(&mut self, hetero_var: &HeteroVar, min_records: usize) -> bool {
        if hetero_var
            .reads
            .iter()
            .filter(|r| self.reads.contains(*r))
            .count()
            >= min_records
        {
            self.reads.extend(hetero_var.reads.clone().into_iter());
            self.data.push(hetero_var.clone());
            self.data.sort_by(|a, b| a.position.cmp(&b.position));
            true
        } else {
            false
        }
    }

    pub fn try_merge_phase(&mut self, phase: &Phase, min_records: usize) -> bool {
        if phase
            .reads
            .iter()
            .filter(|r| self.reads.contains(*r))
            .count()
            >= min_records
        {
            self.reads.extend(phase.reads.clone().into_iter());
            self.data.extend(phase.data.clone().into_iter());
            self.data.sort_by(|a, b| a.position.cmp(&b.position));
            true
        } else {
            false
        }
    }

    pub fn range(&self, bam: &mut IndexedReader) -> Result<(i64, i64, i64, i64)> {
        let mut phase = self.clone();
        phase.sort_dedup();
        let data = &self.data;

        let left = data.first().unwrap();
        let right = data.last().unwrap();

        let left_records = pileup::records_at_base(bam, &left.chr, left.position, true)?;
        let right_records = pileup::records_at_base(bam, &right.chr, right.position, true)?;

        let left_starts: Vec<_> = left_records
            .iter()
            .filter(|(_, b)| *b == left.base)
            .map(|(r, _)| r.pos())
            .collect();
        let right_ends: Vec<_> = right_records
            .iter()
            .filter(|(_, b)| *b == right.base)
            .map(|(r, _)| r.pos() + r.seq_len() as i64)
            .collect();

        let min = left_starts.iter().min();
        let min_cov = left_starts.iter().max();
        let max_cov = right_ends.iter().min();
        let max = right_ends.iter().max();

        if min.is_some() && min_cov.is_some() && max.is_some() && max_cov.is_some() {
            Ok((
                *min.unwrap(),
                *min_cov.unwrap(),
                *max_cov.unwrap(),
                *max.unwrap(),
            ))
        } else {
            Err(anyhow!("problem"))
        }
    }

    pub fn mean_vaf(&self) -> f64 {
        let vafs: Vec<_> = self.data.iter().map(|s| s.vaf).collect();
        vafs.iter().sum::<f64>() / vafs.len() as f64
    }

    pub fn sort_dedup(&mut self) {
        self.data.sort_by(|a, b| a.position.cmp(&b.position));
        self.data
            .dedup_by(|a, b| a.position == b.position && a.chr == b.chr && a.base == b.base);
    }

    pub fn bed_string(&self, bam: &mut IndexedReader) -> Result<String> {
        let first = self.data.first().unwrap();
        let (min, min_cov, max_cov, max) = self.range(bam)?;
        Ok(vec![
            first.chr.to_string(),
            min_cov.to_string(),
            max_cov.to_string(),
            self.mean_vaf().to_string(),
        ]
        .join("\t"))
    }
}

impl Ord for Phase {
    fn cmp(&self, other: &Self) -> Ordering {
        let s = self.data.first().unwrap();
        let other = other.data.first().unwrap();
        s.position.cmp(&other.position)
    }
}

impl PartialOrd for Phase {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for Phase {
    fn eq(&self, other: &Self) -> bool {
        self.reads == other.reads
        // self.data.iter().filter(|s| other.data.contains(s)).count() == other.data.len()
    }
}

impl Eq for Phase {}
