//! Mathematical formulas and traits for prcessing [`Cgats`] [`DataPoint`]s.

use crate::*;
use DataPoint::*;

/// Trait to add values that may or may not be compatible.
/// Returns [`None`] if either value is not a type that can add.
pub trait PartialAdd<Rhs: Sized>: Sized {
    /// The output of a successfull add
    type Output;
    /// Try to add the two values
    fn partial_add(&self, rhs: &Rhs) -> Option<Self::Output>;
}

/// Trait to add values that may or may not be compatible and return [`Self`] if unable to add.
trait AddOrSelf: Sized + Clone + PartialAdd<Self, Output=Self> {
    fn add_or_self(&self, rhs: &Self) -> Self
    {
        self.partial_add(rhs).unwrap_or_else(|| self.clone())
    }
}

impl<T> AddOrSelf for T
where
    T: Sized + Clone + PartialAdd<Self, Output=Self>
{}

/// Trait to sum an iterator of values that may or may not be compatible
pub trait PartialSum<'a, T>: 'a + PartialAdd<Self, Output=T> {
    /// Try to sum the values in the iterator
    fn partial_sum<I: IntoIterator<Item=&'a Self>>(iter: I) -> Option<Self::Output>;
}

impl<'a, T: 'a> PartialSum<'a, Self> for T
where
    T: Clone + PartialAdd<Self, Output=Self>
{
    fn partial_sum<I: IntoIterator<Item=&'a Self>>(iter: I) -> Option<Self::Output> {
        let mut iter = iter.into_iter();
        let mut val = iter.next()?.clone();

        for next in iter {
            val = val.partial_add(next)?;
        }

        Some(val)
    }
}

/// Front end method for the [`PartialSum`] trait
pub trait IterSum<'a, T: 'a + PartialSum<'a, T>>: Sized + IntoIterator<Item=&'a T> {
    /// Try to sum the values in the iterator
    fn iter_sum(self) -> Option<T> {
        PartialSum::partial_sum(self)
    }
}

impl<'a, T, I> IterSum<'a, T> for I
where
    I: Sized + IntoIterator<Item=&'a T>,
    T: PartialSum<'a, T>,
{}

/// Trait to divide values that may not be compatible.
pub trait PartialDiv<Rhs: Sized>: Sized {
    /// The output of a successful division
    type Output;

    /// Divides the values. Returns [`None`] if either value is not a type that can divide.
    fn partial_div(&self, rhs: &Rhs) -> Option<Self::Output>;
}

/// Trait to average values in an iterator that may or may not be compatible
pub trait PartialAvg<'a, T: 'a + PartialSum<'a, T> + PartialDiv<usize, Output=T>>: IterSum<'a, T> {
    /// Try to average the values
    fn partial_avg(self) -> Option<T> {
        let vec = self.into_iter().collect::<Vec<_>>();
        let len = vec.len();
        let sum = vec.iter_sum()?;
        sum.partial_div(&len)
    }
}

impl<'a, T, I> PartialAvg<'a, T> for I
where
    I: IterSum<'a, T>,
    T: PartialSum<'a, T> + PartialDiv<usize, Output=T>,
{}

impl Cgats {
    fn can_add(&self, other: &Self) -> bool {
        self.data_format == other.data_format &&
            self.len() == other.len()
    }
}

impl PartialAdd<Self> for Cgats {
    type Output = Self;
    fn partial_add(&self, rhs: &Self) -> Option<Self::Output> {
        if self.can_add(rhs) {
            let mut sum = self.clone();
            for (lhs, rhs) in sum.iter_mut().zip(rhs.iter()) {
                *lhs = lhs.add_or_self(rhs);
            }
            Some(sum)
        } else {
            None
        }

    }
}

impl Field {
    fn can_avg(&self) -> bool {
        // A field can be averaged unless it is these types
        !matches!(self, SAMPLE_ID | SAMPLE_NAME | BLANK | Other(_))
    }
}

impl PartialDiv<usize> for Cgats {
    type Output = Cgats;
    fn partial_div(&self, rhs: &usize) -> Option<Self::Output> {
        let mut quotient = self.clone();

        for (field, dp) in quotient.iter_mut_with_fields() {
            if field.can_avg() {
                *dp =  dp.partial_div(rhs)?;
            }
        }

        Some(quotient)
    }
}

impl PartialAdd<Self> for DataPoint {
    type Output = Self;
    fn partial_add(&self, rhs: &Self) -> Option<Self::Output> {
        match (self, rhs) {
            (Int(x), Int(y)) => Some(Int(x+y)),
            (Float(x), Float(y)) => Some(Float(x+y)),
            (Int(_), Float(_)) => self.to_float().expect("int to float").partial_add(rhs),
            (Float(_), Int(_)) => self.partial_add(&rhs.to_float().expect("int to float")),
            (Alpha(_), _) | (_, Alpha(_))  => {
                self.to_int().or_else(|_| self.to_float()).ok()?
                    .partial_add(&rhs.to_int().or_else(|_| rhs.to_float()).ok()?)
            }
        }
    }
}

impl PartialAdd<Self> for f32 {
    type Output = Self;
    fn partial_add(&self, rhs: &Self) -> Option<Self::Output> {
        Some(self + rhs)
    }
}

impl PartialDiv<usize> for DataPoint {
    type Output = DataPoint;
    fn partial_div(&self, rhs: &usize) -> Option<Self::Output> {
        let rhs = *rhs as f32;
        match self {
            Int(x) => Some(Float(*x as f32 / rhs)),
            Float(x) => Some(Float(*x / rhs)),
            Alpha(x) => Some(Float(x.parse::<f32>().ok()? / rhs)),
        }
    }
}

impl PartialDiv<Self> for f32 {
    type Output = Self;
    fn partial_div(&self, rhs: &Self) -> Option<Self::Output> {
        Some(self / rhs)
    }
}

impl PartialDiv<usize> for f32 {
    type Output = Self;
    fn partial_div(&self, rhs: &usize) -> Option<Self::Output> {
        Some(self / *rhs as Self)
    }
}

#[test]
fn partial_add() {
    assert_eq!(Int(2).partial_add(&Int(3)), Some(Int(5)));
    assert_eq!(Float(2.0).partial_add(&Float(3.0)), Some(Float(5.0)));
    assert_eq!(Int(2).partial_add(&Float(3.0)), Some(Float(5.0)));
    assert_eq!(Float(2.0).partial_add(&Int(3)), Some(Float(5.0)));
    assert_eq!(Alpha("2".into()).partial_add(&Alpha("3".into())), Some(Int(5)));
    assert_eq!(Alpha("2".into()).partial_add(&Alpha("3.0".into())), Some(Float(5.0)));
    assert_eq!(Alpha("2.0".into()).partial_add(&Alpha("3".into())), Some(Float(5.0)));
    assert_eq!(Alpha("2.0".into()).partial_add(&Int(3)), Some(Float(5.0)));
    assert_eq!(Alpha("2.0".into()).partial_add(&Float(3.0)), Some(Float(5.0)));
    assert_eq!(Int(2).partial_add(&Alpha("3.0".into())), Some(Float(5.0)));
    assert_eq!(Float(2.0).partial_add(&Alpha("3".into())), Some(Float(5.0)));
    assert_eq!(Alpha("x".into()).partial_add(&Int(3)), None);
    assert_eq!(Int(2).partial_add(&Alpha("y".into())), None);
    assert_eq!(Alpha("x".into()).partial_add(&Float(3.0)), None);
    assert_eq!(Float(2.0).partial_add(&Alpha("y".into())), None);
    assert_eq!(Int(2).partial_add(&Alpha("3.0".into())), Some(Float(5.0)));
    assert_eq!(Float(2.0).partial_add(&Alpha("3".into())), Some(Float(5.0)));
}

#[test]
fn partial_sum() {
    let vec = vec![Int(2), Int(3), Int(4)];
    assert_eq!(vec.iter_sum(), Some(Int(9)));
    let vec = vec![Float(2.0), Int(3), Int(4)];
    assert_eq!(vec.iter_sum(), Some(Float(9.0)));
    let vec = vec![Alpha("2.0".into()), Int(3), Int(4)];
    assert_eq!(vec.iter_sum(), Some(Float(9.0)));
    let vec = vec![Alpha("x".into()), Int(3), Int(4)];
    assert_eq!(vec.iter_sum(), None);
}

#[test]
fn partial_div() {
    assert_eq!(Int(6).partial_div(&3), Some(Float(2.0)));
    assert_eq!(Float(6.0).partial_div(&3), Some(Float(2.0)));
    assert_eq!(Alpha("6.0".into()).partial_div(&3), Some(Float(2.0)));
}

#[test]
fn partial_avg() {
    let vec = vec![Int(2), Int(4), Int(6), Int(8)];
    assert_eq!(vec.iter().partial_avg(), Some(Float(5.0)));
    let vec = vec![Int(2), Float(4.0), Alpha("6".into()), Alpha("8".into())];
    assert_eq!(vec.iter().partial_avg(), Some(Float(5.0)));
    let vec = vec![Int(2), Float(4.0), Alpha("6".into()), Alpha("x".into())];
    assert_eq!(vec.iter().partial_avg(), None);
}
#[test]
fn integrated_avg() {
    let cgats: Cgats =
    "CGATS.17
    BEGIN_DATA_FORMAT
    SampleID	RGB_R	RGB_G	RGB_B
    END_DATA_FORMAT
    BEGIN_DATA
    A1	0	1	2
    A2	126	127	128
    A3	253	254	255
    END_DATA"
    .parse().unwrap();

    assert_eq!(cgats.get_col(0).partial_avg(), None);
    assert_eq!(cgats.get_col(1).partial_avg(), Some(Float(126.333_336)));
    assert_eq!(cgats.get_col(2).partial_avg(), Some(Float(127.333_336)));
    assert_eq!(cgats.get_col(3).partial_avg(), Some(Float(128.333_33)));
    assert_eq!(cgats.get_row(0).unwrap().partial_avg(), None);
    assert_eq!(cgats.get_row(0).unwrap().skip(1).partial_avg(), Some(Float(1.0)));
    assert_eq!(cgats.get_row(1).unwrap().skip(1).partial_avg(), Some(Float(127.0)));
    assert_eq!(cgats.get_row(2).unwrap().skip(1).partial_avg(), Some(Float(254.0)));

    let vec = vec![Float(1.23), Float(4.56), Float(9.87)];
    assert_eq!(vec.partial_avg(), Some(Float(5.22)));

    let cgats2: Cgats =
    "CGATS.17
    BEGIN_DATA_FORMAT
    SampleID	RGB_R	RGB_G	RGB_B
    END_DATA_FORMAT
    BEGIN_DATA
    A1	0	1	2
    A2	128	127	126
    A3	255	254	253
    END_DATA"
    .parse().unwrap();

    let vec: Vec<Cgats> = vec![cgats, cgats2];
    let avg: Cgats = vec.partial_avg().unwrap();
    
    let expected: Cgats =
    "CGATS.17
    BEGIN_DATA_FORMAT
    SampleID	RGB_R	RGB_G	RGB_B
    END_DATA_FORMAT
    BEGIN_DATA
    A1	0	1	2
    A2	127	127	127
    A3	254	254	254
    END_DATA"
    .parse().unwrap();

    println!("{}", &avg);
    println!("{}", &expected);

    assert_eq!(avg, expected);
}