<commit_before>use std::cmp::{max, min};
use std::collections::HashMap;

use bio::utils::{self, Interval, IntervalError};
use bio::utils::Strand;

use {Coord, Error};
use self::ExonFeatureKind::*;


macro_rules! impl_common {
    ($struct_ty:ty) => (

        impl $struct_ty {

            pub fn seq_name(&self) -> &str {
                self.seq_name.as_str()
            }

            pub fn interval(&self) -> &Interval<u64> {
                &self.interval
            }

            pub fn strand(&self) -> &Strand {
                &self.strand
            }

            #[inline]
            pub fn span(&self) -> u64 {
                self.interval().end - self.interval().start
            }
        }

    );
}

#[derive(Debug, Clone, PartialEq)]
pub struct Feature<K: FeatureKind> {
    pub interval: Interval<u64>,
    pub kind: K,
}

impl<K: FeatureKind> Feature<K> {

    #[inline]
    fn span(&self) -> u64 {
        self.interval.end - self.interval.start
    }
}

pub trait FeatureKind {}

#[derive(Debug, Clone, PartialEq)]
pub enum ExonFeatureKind {
    UTR,
    UTR5,
    UTR3,
    CDS { frame: Option<u64> },
    StartCodon { frame: Option<u64> },
    StopCodon { frame: Option<u64> },
    Any(String),
}

impl FeatureKind for ExonFeatureKind {}

pub type ExonFeature = Feature<ExonFeatureKind>;

#[derive(Debug, Clone, PartialEq)]
pub enum TranscriptFeatureKind {
    Intron,
    Any(String),
}

impl FeatureKind for TranscriptFeatureKind {}

pub type TranscriptFeature = Feature<TranscriptFeatureKind>;

#[derive(Debug, Clone, PartialEq)]
pub struct GeneFeatureKind(String);

impl FeatureKind for GeneFeatureKind {}

pub type GeneFeature = Feature<GeneFeatureKind>;

#[derive(Debug)]
pub struct Exon {
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    pub id: Option<String>,
    pub attributes: HashMap<String, String>,
    pub features: Vec<ExonFeature>,
}

impl_common!(Exon);

pub struct EBuilder {
    seq_name: String,
    start: u64,
    end: u64,
    strand: Option<Strand>,
    strand_char: Option<char>,
    pub id: Option<String>,
    pub attributes: HashMap<String, String>,
    pub features: Vec<ExonFeature>,
}

impl EBuilder {

    pub fn new<T>(seq_name: T, start: u64, end: u64) -> Self
        where T: Into<String>
    {
        EBuilder {
            seq_name: seq_name.into(),
            start: start,
            end: end,
            strand: None,
            strand_char: None,
            id: None,
            attributes: HashMap::new(),
            features: Vec::new(),
        }
    }

    pub fn strand(mut self, strand: Strand) -> Self {
        self.strand = Some(strand);
        self
    }

    pub fn strand_char(mut self, strand_char: char) -> Self {
        self.strand_char = Some(strand_char);
        self
    }

    pub fn id<T>(mut self, id: T) -> Self
        where T: Into<String>
    {
        self.id = Some(id.into());
        self
    }

    pub fn attribute<K, V>(mut self, key: K, value: V) -> Self
        where K: Into<String>, V: Into<String>
    {
        self.attributes.insert(key.into(), value.into());
        self
    }

    pub fn attributes(mut self, attributes: HashMap<String, String>) -> Self {
        self.attributes = attributes;
        self
    }

    pub fn feature(mut self, feature: ExonFeature) -> Self {
        self.features.push(feature);
        self
    }

    pub fn features(mut self, features: Vec<ExonFeature>) -> Self {
        self.features = features;
        self
    }

    pub fn build(self) -> Result<Exon, Error> {
        let interval = coord_to_interval(self.start, self.end)
            .map_err(Error::Feature)?;
        let strand = resolve_strand_input(self.strand, self.strand_char)
            .map_err(Error::Feature)?;
        let feature = Exon {
            seq_name: self.seq_name,
            interval: interval,
            strand: strand,
            id: self.id,
            attributes: self.attributes,
            features: self.features,
        };
        Ok(feature)
    }
}

#[derive(Debug)]
pub struct Transcript {
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    pub id: Option<String>,
    pub attributes: HashMap<String, String>,
    exons: Vec<Exon>,
}

impl_common!(Transcript);

impl Transcript {

    pub fn exons(&self) -> &[Exon] {
        self.exons.as_slice()
    }
}

pub struct TBuilder {
    seq_name: String,
    start: u64,
    end: u64,
    strand: Option<Strand>,
    strand_char: Option<char>,
    id: Option<String>,
    attributes: HashMap<String, String>,
    // Input can be a vector of pre-made features ...
    exons: Option<Vec<Exon>>,
    // Or exon coordinates, possibly coupled with cds coord
    // NOTE: Can we instead of using Vec<_> here keep it as an unconsumed iterator?
    exon_coords: Option<Vec<Coord<u64>>>,
    coding_coord: Option<Coord<u64>>,
    coding_incl_stop: bool,
}<commit_msg>Add Transcript fn to return coding coords<commit_after>use std::cmp::{max, min};
use std::collections::HashMap;

use bio::utils::{self, Interval, IntervalError};
use bio::utils::Strand;

use {Coord, Error};
use self::ExonFeatureKind::*;


macro_rules! impl_common {
    ($struct_ty:ty) => (

        impl $struct_ty {

            pub fn seq_name(&self) -> &str {
                self.seq_name.as_str()
            }

            pub fn interval(&self) -> &Interval<u64> {
                &self.interval
            }

            pub fn strand(&self) -> &Strand {
                &self.strand
            }

            #[inline]
            pub fn span(&self) -> u64 {
                self.interval().end - self.interval().start
            }
        }

    );
}

#[derive(Debug, Clone, PartialEq)]
pub struct Feature<K: FeatureKind> {
    pub interval: Interval<u64>,
    pub kind: K,
}

impl<K: FeatureKind> Feature<K> {

    #[inline]
    fn span(&self) -> u64 {
        self.interval.end - self.interval.start
    }
}

pub trait FeatureKind {}

#[derive(Debug, Clone, PartialEq)]
pub enum ExonFeatureKind {
    UTR,
    UTR5,
    UTR3,
    CDS { frame: Option<u64> },
    StartCodon { frame: Option<u64> },
    StopCodon { frame: Option<u64> },
    Any(String),
}

impl FeatureKind for ExonFeatureKind {}

pub type ExonFeature = Feature<ExonFeatureKind>;

#[derive(Debug, Clone, PartialEq)]
pub enum TranscriptFeatureKind {
    Intron,
    Any(String),
}

impl FeatureKind for TranscriptFeatureKind {}

pub type TranscriptFeature = Feature<TranscriptFeatureKind>;

#[derive(Debug, Clone, PartialEq)]
pub struct GeneFeatureKind(String);

impl FeatureKind for GeneFeatureKind {}

pub type GeneFeature = Feature<GeneFeatureKind>;

#[derive(Debug)]
pub struct Exon {
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    pub id: Option<String>,
    pub attributes: HashMap<String, String>,
    pub features: Vec<ExonFeature>,
}

impl_common!(Exon);

pub struct EBuilder {
    seq_name: String,
    start: u64,
    end: u64,
    strand: Option<Strand>,
    strand_char: Option<char>,
    pub id: Option<String>,
    pub attributes: HashMap<String, String>,
    pub features: Vec<ExonFeature>,
}

impl EBuilder {

    pub fn new<T>(seq_name: T, start: u64, end: u64) -> Self
        where T: Into<String>
    {
        EBuilder {
            seq_name: seq_name.into(),
            start: start,
            end: end,
            strand: None,
            strand_char: None,
            id: None,
            attributes: HashMap::new(),
            features: Vec::new(),
        }
    }

    pub fn strand(mut self, strand: Strand) -> Self {
        self.strand = Some(strand);
        self
    }

    pub fn strand_char(mut self, strand_char: char) -> Self {
        self.strand_char = Some(strand_char);
        self
    }

    pub fn id<T>(mut self, id: T) -> Self
        where T: Into<String>
    {
        self.id = Some(id.into());
        self
    }

    pub fn attribute<K, V>(mut self, key: K, value: V) -> Self
        where K: Into<String>, V: Into<String>
    {
        self.attributes.insert(key.into(), value.into());
        self
    }

    pub fn attributes(mut self, attributes: HashMap<String, String>) -> Self {
        self.attributes = attributes;
        self
    }

    pub fn feature(mut self, feature: ExonFeature) -> Self {
        self.features.push(feature);
        self
    }

    pub fn features(mut self, features: Vec<ExonFeature>) -> Self {
        self.features = features;
        self
    }

    pub fn build(self) -> Result<Exon, Error> {
        let interval = coord_to_interval(self.start, self.end)
            .map_err(Error::Feature)?;
        let strand = resolve_strand_input(self.strand, self.strand_char)
            .map_err(Error::Feature)?;
        let feature = Exon {
            seq_name: self.seq_name,
            interval: interval,
            strand: strand,
            id: self.id,
            attributes: self.attributes,
            features: self.features,
        };
        Ok(feature)
    }
}

#[derive(Debug)]
pub struct Transcript {
    seq_name: String,
    interval: Interval<u64>,
    strand: Strand,
    pub id: Option<String>,
    pub attributes: HashMap<String, String>,
    exons: Vec<Exon>,
}

impl_common!(Transcript);

impl Transcript {

    pub fn exons(&self) -> &[Exon] {
        self.exons.as_slice()
    }

    pub fn coding_coord(&self, incl_stop: bool) -> Option<Coord<u64>> {
        let start = self.coding_start_coord(incl_stop);
        let end = self.coding_end_coord(incl_stop);
        match (start, end) {
            (Some(s), Some(e)) => Some((s, e)),
            _ => None,
        }
    }

    fn coding_start_coord(&self, incl_stop: bool) -> Option<u64> {
        match &self.strand {
            &Strand::Forward => {
                for exon in self.exons.iter() {
                    for fx in exon.features.iter() {
                        if let StartCodon { .. } = fx.kind {
                            return Some(fx.interval.start);
                        }
                    }
                }
                None
            },
            &Strand::Reverse => {
                let mut codon_rem = if incl_stop { 0 } else { 3 };
                for exon in self.exons.iter() {
                    for fx in exon.features.iter() {
                        if let StopCodon { .. } = fx.kind {
                            if incl_stop {
                                return Some(fx.interval.start);
                            }
                            codon_rem -= fx.span();
                            if codon_rem == 0 {
                                return Some(fx.interval.end);
                            }
                        }
                    }
                }
                None
            },
            &Strand::Unknown if incl_stop => {
                for exon in self.exons.iter() {
                    for fx in exon.features.iter() {
                        if let CDS { .. } = fx.kind {
                            return Some(fx.interval.start)
                        }
                    }
                }
                None
            },
            _ => None
        }
    }

    fn coding_end_coord(&self, incl_stop: bool) -> Option<u64> {
        match &self.strand {
            &Strand::Forward => {
                let mut codon_rem = if incl_stop { 0 } else { 3 };
                for exon in self.exons.iter().rev() {
                    for fx in exon.features.iter().rev() {
                        if let StopCodon { .. } = fx.kind {
                            if incl_stop {
                                return Some(fx.interval.end);
                            }
                            codon_rem -= fx.span();
                            if codon_rem == 0 {
                                return Some(fx.interval.start);
                            }
                        }
                    }
                }
                None
            },
            &Strand::Reverse => {
                for exon in self.exons.iter().rev() {
                    for fx in exon.features.iter().rev() {
                        if let StartCodon { .. } = fx.kind {
                            return Some(fx.interval.end);
                        }
                    }
                }
                None
            },
            &Strand::Unknown if incl_stop => {
                for exon in self.exons.iter().rev() {
                    for fx in exon.features.iter().rev() {
                        if let CDS { .. } = fx.kind {
                            return Some(fx.interval.end);
                        }
                    }
                }
                None
            },
            _ => None,
        }
    }

}

pub struct TBuilder {
    seq_name: String,
    start: u64,
    end: u64,
    strand: Option<Strand>,
    strand_char: Option<char>,
    id: Option<String>,
    attributes: HashMap<String, String>,
    // Input can be a vector of pre-made features ...
    exons: Option<Vec<Exon>>,
    // Or exon coordinates, possibly coupled with cds coord
    // NOTE: Can we instead of using Vec<_> here keep it as an unconsumed iterator?
    exon_coords: Option<Vec<Coord<u64>>>,
    coding_coord: Option<Coord<u64>>,
    coding_incl_stop: bool,
}
{'fraction_non_alphanumeric': 0.11868231046931407, 'fraction_numerical': 0.00631768953068592, 'mean_word_length': 1.663661858974359, 'pattern_counts': {'":': 0, '<': 99, '<?xml version=': 0, '>': 144, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '781623', 'n_tokens_mistral': 4113, 'n_tokens_neox': 3908, 'n_words': 966}