rbop 0.2.0

use core::{cmp::Ordering, convert::TryInto, ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign}};

use alloc::{vec, vec::Vec};
use num_integer::{Roots, Integer};
use num_traits::{FromPrimitive, One, ToPrimitive, Zero};
use rust_decimal::{Decimal, MathematicalOps};

use crate::{decimal_ext::DecimalExtensions, node::unstructured::Serializable, error::MathsError};

#[derive(PartialEq, Eq, Copy, Clone, Debug)]
pub enum Number {
    Decimal(Decimal),
    Rational(i64, i64),
}

impl Number {
    /// Converts this number to a decimal:
    ///   - For `Decimal`, this simply unwraps the variant.
    ///   - For `Rational`, this divides the numerator by the denominator after converting both to
    ///     decimals.
    pub fn to_decimal(&self) -> Decimal {
        match self {
            Number::Decimal(d) => *d,
            Number::Rational(numer, denom)
                => Decimal::from_i64(*numer).unwrap() / Decimal::from_i64(*denom).unwrap(),
        }
    }

    /// Utility function which gets the greatest common denominator of two numbers. 
    fn gcd(a: i64, b: i64) -> i64 {
        if b == 0 {
            return a;
        }

        Self::gcd(b, a % b)
    }

    /// Utility function which gets the lowest common multiple of two numbers.
    fn lcm(a: i64, b: i64) -> i64 {
        (a * b).abs() / Self::gcd(a, b)
    }

    /// Given two `Rational` numbers, returns the same two numbers in the form (self, other), except
    /// that both numbers have the same denominator.
    ///
    /// Panics if either of the numbers is not rational.
    fn to_common_with(self, other: Number) -> (Number, Number) {
        if let (Self::Rational(ln, ld), Self::Rational(rn, rd)) = (self, other) {
            let new_denominator = Self::lcm(ld, rd);
            let ln = (new_denominator / ld) * ln;
            let rn = (new_denominator / rd) * rn;

            (
                Self::Rational(ln, new_denominator),
                Self::Rational(rn, new_denominator),
            )
        } else {
            panic!("both numbers must be rational");
        }
    }

    /// Assumes that this is a `Rational` number and returns its numerator, otherwise panics.
    fn numerator(&self) -> i64 {
        if let Self::Rational(n, _) = self {
            *n
        } else {
            panic!("not rational")
        }
    }

    /// Assumes that this is a `Rational` number and returns its denominator, otherwise panics.
    fn denominator(&self) -> i64 {
        if let Self::Rational(_, d) = self {
            *d
        } else {
            panic!("not rational")
        }
    }

    /// Simplifies this number:
    ///   - For `Decimal`, this normalises the number. This performs lossless conversions such as
    ///     removing trailing zeroes (which can occur in rust_decimal), and converting -0 to 0.
    ///   - For `Rational`, this divides the numerator and denominator by their GCD. Also ensures
    ///     that any negative sign is on the numerator, not the denominator.
    pub fn simplify(&self) -> Number {
        match self {
            Self::Decimal(d) => Self::Decimal(d.normalize()),

            Self::Rational(numer, denom) => {
                let sign = match (*numer < 0, *denom < 0) {
                    (false, false) => 1, // Neither negative
                    (true, false) | (false, true) => -1, // One negative
                    (true, true) => 1, // Both negative, cancels out
                };

                let (numer, denom) = (numer.abs(), denom.abs());

                let gcd = Self::gcd(numer, denom);
                Self::Rational(sign * (numer / gcd), denom / gcd)
            }
        }
    }

    /// Returns the reciprocal of this number.
    pub fn reciprocal(&self) -> Number {
        match self {
            Self::Decimal(d) => Self::Decimal(Decimal::one() / d),
            Self::Rational(numer, denom) => Self::Rational(*denom, *numer),
        }
    }

    /// If this is a whole number, returns it. Otherwise returns None.
    pub fn to_whole(&self) -> Option<i64> {
        match self {
            Self::Decimal(d)
                => if d.is_whole() { d.floor().to_i64() } else { None },
            Self::Rational(numer, denom)
                => if numer % denom == 0 { Some(numer / denom) } else { None },
        }
    }

    /// Raises this number to an integer power.
    pub fn powi(&self, exp: i64) -> Number {
        let mut n = *self;

        // Repeatedly multiply 
        for _ in 1..exp.abs() {
            n = n * *self;
        }
        
        // Reciprocal for negative powers
        if exp < 0 {
            n.reciprocal()
        } else {
            n
        }
    }

    /// Adds this number to another number, or returns an error if an overflow occurs.
    pub fn checked_add(&self, other: Number) -> Result<Number, MathsError> {
        if let (l@Self::Rational(_, _), r@Self::Rational(_, _)) = (self, other) {
            let (l, r) = l.to_common_with(r);

            Ok(Number::Rational(
                l.numerator().checked_add(r.numerator()).ok_or(MathsError::Overflow)?,
                l.denominator(),
            ).simplify())
        } else {
            Ok(Number::Decimal(
                self.to_decimal().checked_add(other.to_decimal()).ok_or(MathsError::Overflow)?,
            ))
        }
    }

    /// Subtracts one number from another number, or returns an error if an overflow occurs.
    pub fn checked_sub(&self, other: Number) -> Result<Number, MathsError> {
        self.checked_add(-other)
    }

    /// Multiplies this number with another number, or returns an error if an overflow occurs.
    pub fn checked_mul(&self, other: Number) -> Result<Number, MathsError> {
        if let (Self::Rational(ln, ld), Self::Rational(rn, rd)) = (self, other) {
            Ok(Number::Rational(
                ln.checked_mul(rn).ok_or(MathsError::Overflow)?,
                ld.checked_mul(rd).ok_or(MathsError::Overflow)?,
            ).simplify())
        } else {
            Ok(Number::Decimal(
                self.to_decimal().checked_mul(other.to_decimal()).ok_or(MathsError::Overflow)?,
            ).simplify())
        }
    }

    /// Divides this number by another number, or returns an error if the divisor is zero.
    pub fn checked_div(&self, other: Number) -> Result<Number, MathsError> {
        if other.is_zero() {
            Err(MathsError::DivisionByZero)
        } else {
            Ok(*self / other)
        }
    }

    /// Raises this number to the power of another number.
    pub fn checked_pow(&self, power: Number) -> Result<Number, MathsError> {
        // If both power and base are rational, we can get a bit more accuracy by breaking it down
        if let (Self::Rational(bn, bd), Self::Rational(pn, pd)) = (self, power) {
            // Can only keep as rational if (power denominator)th root of both base numerator and
            // denominator are integers
            if (bn.is_negative() && pd.is_even()) || (bd.is_negative() && pd.is_even()) {
                return Err(MathsError::Imaginary)
            }

            // TODO: handle panics in `nth_root`
            let bn_pd_nth_root = bn.nth_root(pd.try_into().map_err(|_| MathsError::Overflow)?);
            let bd_pd_nth_root = bd.nth_root(pd.try_into().map_err(|_| MathsError::Overflow)?);
            if bn_pd_nth_root.pow(pd.abs().try_into().map_err(|_| MathsError::Overflow)?) == *bn
               && bd_pd_nth_root.pow(pd.abs().try_into().map_err(|_| MathsError::Overflow)?) == *bd {

                let mut result = Number::Rational(
                    bn_pd_nth_root.pow(pn.abs().try_into().map_err(|_| MathsError::Overflow)?), 
                    bd_pd_nth_root.pow(pn.abs().try_into().map_err(|_| MathsError::Overflow)?), 
                );

                if pn < 0 {
                    result = result.reciprocal();
                }

                return Ok(result)
            }
        }

        Ok(Number::Decimal(
            self.to_decimal().checked_powd(power.to_decimal()).ok_or(MathsError::Overflow)?
        ))
    }
}

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

impl Ord for Number {
    fn cmp(&self, other: &Self) -> Ordering {
        self.to_decimal().cmp(&other.to_decimal())
    }
}

impl From<Decimal> for Number {
    fn from(d: Decimal) -> Self {
        Self::Decimal(d)
    }
}

impl From<i64> for Number {
    fn from(i: i64) -> Self {
        Self::Rational(i, 1)
    }
}

impl Neg for Number {
    type Output = Self;

    fn neg(self) -> Self::Output {
        match self {
            Self::Rational(n, d) => Number::Rational(-n, d).simplify(),
            Self::Decimal(d) => Self::Decimal(-d),
        }
    }
}

impl Add for Number {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        self.checked_add(rhs).unwrap()
    }
}

impl AddAssign for Number {
    fn add_assign(&mut self, rhs: Self) { *self = *self + rhs; }
}

impl Sub for Number {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        self + -rhs
    }
}

impl SubAssign for Number {
    fn sub_assign(&mut self, rhs: Self) { *self = *self - rhs; }
}

impl Mul for Number {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        self.checked_mul(rhs).unwrap()
    }
}

impl MulAssign for Number {
    fn mul_assign(&mut self, rhs: Self) { *self = *self * rhs; }
}

impl Div for Number {
    type Output = Self;
    
    #[allow(clippy::suspicious_arithmetic_impl)]
    fn div(self, rhs: Self) -> Self::Output {
        self * rhs.reciprocal()
    }
}

impl DivAssign for Number {
    fn div_assign(&mut self, rhs: Self) { *self = *self / rhs; }
}

impl Zero for Number {
    fn zero() -> Self {
        Self::Rational(0, 1)
    }

    fn is_zero(&self) -> bool {
        match *self {
            Self::Decimal(d) => d.is_zero(),
            Self::Rational(n, _) => n.is_zero(),
        }
    }
}

impl One for Number {
    fn one() -> Self {
        Self::Rational(1, 1)
    }

    fn is_one(&self) -> bool {
        match *self {
            Self::Decimal(d) => d.is_one(),
            Self::Rational(n, d) => n == d,
        }
    }
}

impl Serializable for Number {
    fn serialize(&self) -> Vec<u8> {
        match self {
            Number::Decimal(d) => {
                let mut result = vec![1];
                result.append(&mut d.serialize().to_vec());
                result
            }

            Self::Rational(numer, denom) => {
                let mut result = vec![2];
                result.append(&mut numer.to_ne_bytes().to_vec());
                result.append(&mut denom.to_ne_bytes().to_vec());
                result
            }
        }
    }

    fn deserialize(bytes: &mut dyn Iterator<Item = u8>) -> Option<Self> {
        let first_byte = bytes.next()?;
        match first_byte {
            1 => {
                Some(Number::Decimal(Decimal::deserialize(
                    bytes.take(16).collect::<Vec<_>>().try_into().ok()?
                )))
            }

            2 => {
                let numer: [u8; 8] = bytes.take(8).collect::<Vec<_>>().try_into().ok()?;
                let denom: [u8; 8] = bytes.take(8).collect::<Vec<_>>().try_into().ok()?;
                Some(Number::Rational(
                    i64::from_ne_bytes(numer),
                    i64::from_ne_bytes(denom),
                ))
            }

            _ => None
        }
    }
}