use std::{cmp::Ordering, fmt::Display, ops::Range, str::FromStr};

use anyhow::Context;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};

// 0-based
#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize, Hash)]
pub struct GenomePosition {
    pub contig: u8,
    pub position: u32,
}

impl GenomePosition {
    pub fn contig(&self) -> String {
        num_to_string(self.contig)
    }
}

impl Display for GenomePosition {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}:{}", self.contig(), self.position)
    }
}

impl FromStr for GenomePosition {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        let (c, p) = s.split_once(":").ok_or(anyhow::anyhow!("Can't split {s}"))?;
        Ok(Self {
            contig: contig_to_num(c),
            position: p.parse()?,
        })
    }
}

impl TryFrom<(&str, &str)> for GenomePosition {
    type Error = anyhow::Error;

    fn try_from((c, p): (&str, &str)) -> anyhow::Result<Self> {
        Ok(GenomePosition {
            contig: contig_to_num(c),
            position: p.parse().context(format!("Can't parse {p} into u32."))?,
        })
    }
}

impl From<VcfPosition> for GenomePosition {
    fn from(value: VcfPosition) -> Self {
        Self {
            contig: value.contig,
            position: value.position - 1,
        }
    }
}

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

impl From<GenomePosition> for (String, String) {
    fn from(val: GenomePosition) -> Self {
        (num_to_string(val.contig), val.position.to_string())
    }
}

impl Ord for GenomePosition {
    fn cmp(&self, other: &Self) -> Ordering {
        match self.contig.cmp(&other.contig) {
            Ordering::Equal => self.position.cmp(&other.position),
            other => other,
        }
    }
}

pub fn contig_to_num(contig: &str) -> u8 {
    match contig {
        "chrX" => 23,
        "chrY" => 24,
        "chrM" => 25,
        c => c.trim_start_matches("chr").parse().unwrap_or(u8::MAX),
    }
}

pub fn num_to_string(num: u8) -> String {
    match num {
        23 => "chrX".to_string(),
        24 => "chrY".to_string(),
        25 => "chrM".to_string(),
        c => format!("chr{c}"),
    }
}

// 1-based
#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
pub struct VcfPosition {
    pub contig: u8,
    pub position: u32,
}

impl TryFrom<(&str, &str)> for VcfPosition {
    type Error = anyhow::Error;

    fn try_from((c, p): (&str, &str)) -> anyhow::Result<Self> {
        Ok(VcfPosition {
            contig: contig_to_num(c),
            position: p.parse().context(format!("Can't parse {p} into u32."))?,
        })
    }
}

impl From<GenomePosition> for VcfPosition {
    fn from(value: GenomePosition) -> Self {
        Self {
            contig: value.contig,
            position: value.position + 1,
        }
    }
}

impl From<VcfPosition> for (String, String) {
    fn from(val: VcfPosition) -> Self {
        (num_to_string(val.contig), val.position.to_string())
    }
}

#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
pub struct GenomeRange {
    pub contig: u8,
    pub range: Range<u32>,
}

impl TryFrom<(&str, &str, &str)> for GenomeRange {
    type Error = anyhow::Error;

    fn try_from((c, s, e): (&str, &str, &str)) -> anyhow::Result<Self> {
        Ok(Self {
            contig: contig_to_num(c),
            range: Range {
                start: s.parse().context(format!("Can't parse {s} into u32."))?,
                end: e.parse().context(format!("Can't parse {e} into u32."))?,
            },
        })
    }
}

pub fn overlaps_par(positions: &[&GenomePosition], ranges: &[&GenomeRange]) -> Vec<usize> {
    let chunk_size = (positions.len() / rayon::current_num_threads()).max(1);

    positions
        .par_chunks(chunk_size)
        .enumerate()
        .flat_map(|(chunk_idx, chunk)| {
            let mut result = Vec::new();
            let mut range_idx = ranges.partition_point(|r| r.contig < chunk[0].contig);

            for (pos_idx, pos) in chunk.iter().enumerate() {
                while range_idx < ranges.len() && ranges[range_idx].contig < pos.contig {
                    range_idx += 1;
                }

                if range_idx < ranges.len() && ranges[range_idx].contig == pos.contig {
                    while range_idx < ranges.len()
                        && ranges[range_idx].contig == pos.contig
                        && ranges[range_idx].range.end <= pos.position
                    {
                        range_idx += 1;
                    }

                    if range_idx < ranges.len()
                        && ranges[range_idx].contig == pos.contig
                        && ranges[range_idx].range.contains(&pos.position)
                    {
                        result.push(chunk_idx * chunk_size + pos_idx);
                    }
                }
            }
            result
        })
        .collect()
}

pub fn par_overlaps(a: &[impl GetGenomePosition], b: &[impl GetGenomeRange]) -> Vec<usize> {
    let a: Vec<_> = a.iter().map(|e| e.position()).collect();
    let b: Vec<_> = b.iter().map(|e| e.range()).collect();
    overlaps_par(&a, &b)
}

pub trait GetGenomePosition {
    fn position(&self) -> &GenomePosition;
}

pub trait GetGenomeRange {
    fn range(&self) -> &GenomeRange;
}
