#![feature(generic_const_exprs)]
#![feature(adt_const_params)]
#![cfg_attr(not(test), no_std)]

//! A library that lets you validate the dimensions of your quantities at compile time
//!
//! WARNING: Uses the experimental features `generic_const_exprs` and `adt_const_params` which are known to cause bugs
//!
//! ## Example
//!
//! Okay
//!
//! ```
//! use const_units::prefixes::atto; // prefix for 10^-18
//! use const_units::units::meter;
//! use const_units::Quantity; // Represents a number with a dimension
//!
//! // Input attometers, return square attometers
//! fn square_dist(
//!     x: Quantity<f64, { atto(meter) }>,
//!     y: Quantity<f64, { atto(meter) }>,
//! ) -> Quantity<f64, { atto(meter).powi(2) }> {
//!     x.powi::<2>() + y.powi::<2>()
//! }
//!
//! // Input attometers, return attometers
//! fn distance(x: Quantity<f64, { atto(meter) }>, y: Quantity<f64, { atto(meter) }>) -> Quantity<f64, { atto(meter) }> {
//!     square_dist(x, y).powf::<{ (1, 2) }>() // `(1, 2)` represents 1/2
//! }
//! ```
//!
//! Broken
//!
//! ```compile_fail
//! # use const_units::{Quantity, units::{meter, second, DIMENSIONLESS}};
//! fn sum(x: Quantity<f64, { meter }>, y: Quantity<f64, { second }>) -> Quantity<f64, DIMENSIONLESS> {
//!     x + y // You can't add meters and seconds
//! }
//! ```

use core::{
    fmt::Display,
    marker::ConstParamTy,
    ops::{Add, AddAssign, Deref, DerefMut, Div, DivAssign, Mul, MulAssign, Sub, SubAssign},
};

use num::{traits::Pow, NumCast};
use units::DIMENSIONLESS;

pub mod prefixes;
pub mod units;

// https://stackoverflow.com/questions/7407752/integer-nth-root
const fn try_root(v: u128, n: u32) -> Option<u128> {
    // (2..128).map(|v| (u128::MAX as f64).powf(1. / v as f64).ceil() as u128).collect::<Vec<_>>()
    #[rustfmt::skip]
    const MAX_VALS: [u128; 126] = [18446744073709551616, 6981463658332, 4294967296, 50859009, 2642246, 319558, 65536, 19113, 7132, 3184, 1626, 921, 566, 371, 256, 185, 139, 107, 85, 69, 57, 48, 41, 35, 31, 27, 24, 22, 20, 18, 16, 15, 14, 13, 12, 12, 11, 10, 10, 9, 9, 8, 8, 8, 7, 7, 7, 7, 6, 6, 6, 6, 6, 6, 5, 5, 5, 5, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3];

    if n == 1 {
        return Some(v);
    }

    if n == 0 {
        return Some(1);
    }

    if v == 0 || v == 1 {
        return Some(v);
    }

    if n > (MAX_VALS.len() + 1) as u32 {
        return None;
    }

    // Can't shadow constant `min`
    let mut minimum = 2;
    let mut max = MAX_VALS[n as usize - 2];

    loop {
        let mid = (minimum + max) >> 1;

        let res = mid.pow(n);

        if res > v {
            max = mid;
        } else if res < v {
            minimum = mid + 1;
        } else {
            return Some(mid);
        }

        if minimum == max {
            return None;
        }
    }
}

const fn add(lhs: (i32, u32), rhs: (i32, u32)) -> (i32, u32) {
    simplify((
        lhs.0 * (rhs.1 as i32) + rhs.0 * (lhs.1 as i32),
        lhs.1 * rhs.1,
    ))
}

const fn sub(lhs: (i32, u32), mut rhs: (i32, u32)) -> (i32, u32) {
    rhs.0 = -rhs.0;

    add(lhs, rhs)
}

const fn recip(f: (i128, u128)) -> (i128, u128) {
    (f.1 as i128 * f.0.signum(), f.0.abs() as u128)
}

const fn mul(lhs: (i128, u128), rhs: (i128, u128)) -> (i128, u128) {
    simplify_big((lhs.0 * rhs.0, lhs.1 * rhs.1))
}

const fn div(lhs: (i128, u128), rhs: (i128, u128)) -> (i128, u128) {
    mul(lhs, recip(rhs))
}

const fn gcd(a: u32, b: u32) -> u32 {
    if b == 0 {
        a
    } else {
        gcd(b, a % b)
    }
}

const fn gcd_big(a: u128, b: u128) -> u128 {
    if b == 0 {
        a
    } else {
        gcd_big(b, a % b)
    }
}

/// Simplify a fraction that uses 32 bit numbers
pub const fn simplify(v: (i32, u32)) -> (i32, u32) {
    let scale_down = gcd(v.0.abs() as u32, v.1);

    (v.0 / (scale_down as i32), v.1 / scale_down)
}

/// Simplify a fraction that uses 128 bit numbers
pub const fn simplify_big(v: (i128, u128)) -> (i128, u128) {
    let scale_down = gcd_big(v.0.abs() as u128, v.1);

    (v.0 / (scale_down as i128), v.1 / scale_down)
}

/// Represents SI units
///
/// Each field contains the exponent of the corresponding dimension. Exponents are fractions where the first element of the tuple is the numerator and the second is the denominator. Ensure the fractions are fully simplified if you're creating an instance directly.
#[allow(non_snake_case)]
#[derive(PartialEq, Eq, Debug, Clone, Copy)]
pub struct SI {
    /// A scale factor - used for units that are constant multiples of other units; ie. minutes or kilometers
    pub scale: (i128, u128),
    /// Seconds
    pub s: (i32, u32),
    /// Meters
    pub m: (i32, u32),
    /// Kilograms
    pub kg: (i32, u32),
    /// Amperes
    pub A: (i32, u32),
    /// Kelvin
    pub K: (i32, u32),
    /// Moles
    pub mol: (i32, u32),
    /// Candelas
    pub cd: (i32, u32),
}

impl Display for SI {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let mut not_first = false;

        if self.scale.0 != 1
            || self.scale.1 != 1
            || (*self).scale_by(recip(self.scale)) == DIMENSIONLESS
        {
            f.write_str("(")?;
            self.scale.0.fmt(f)?;

            if self.scale.1 != 1 {
                f.write_str("/")?;
                self.scale.1.fmt(f)?;
            }

            f.write_str(")")?;

            not_first = true;
        }

        macro_rules! unit {
            ($name: tt, $signum: expr) => {
                #[allow(unused_assignments)]
                // Implicitly checks for zero too
                if self.$name.0.signum() == $signum {
                    if not_first {
                        f.write_str("⋅")?;
                    }

                    f.write_str(stringify!($name))?;

                    if self.$name.0 != 1 || self.$name.1 != 1 {
                        f.write_str("^")?;
                        self.$name.0.fmt(f)?;

                        if self.$name.1 != 1 {
                            f.write_str("/")?;
                            self.$name.1.fmt(f)?;
                        }
                    }

                    not_first = true;
                }
            };
        }

        // Ensure negative powers come after positive ones

        unit!(s, 1);
        unit!(m, 1);
        unit!(kg, 1);
        unit!(A, 1);
        unit!(K, 1);
        unit!(mol, 1);
        unit!(cd, 1);

        unit!(s, -1);
        unit!(m, -1);
        unit!(kg, -1);
        unit!(A, -1);
        unit!(K, -1);
        unit!(mol, -1);
        unit!(cd, -1);

        Ok(())
    }
}

impl SI {
    /// Multiply the scale factor by a given constant
    ///
    /// ```
    /// # use const_units::units::{second, minute};
    /// assert_eq!(second.scale_by((60, 1)), minute)
    /// ```
    pub const fn scale_by(mut self, scale: (i128, u128)) -> Self {
        self.scale = mul(scale, self.scale);
        self
    }

    /// Multiply SI units
    ///
    /// ```
    /// # use const_units::units::{newton, meter, joule};
    /// assert_eq!(newton.mul(meter), joule)
    /// ```
    pub const fn mul(self, rhs: Self) -> Self {
        SI {
            scale: mul(self.scale, rhs.scale),
            s: add(self.s, rhs.s),
            m: add(self.m, rhs.m),
            kg: add(self.kg, rhs.kg),
            A: add(self.A, rhs.A),
            K: add(self.K, rhs.K),
            mol: add(self.mol, rhs.mol),
            cd: add(self.cd, rhs.cd),
        }
    }

    /// Divide SI units
    ///
    /// ```
    /// # use const_units::units::{joule, second, watt};
    /// assert_eq!(joule.div(second), watt);
    /// ```
    pub const fn div(self, rhs: Self) -> Self {
        SI {
            scale: div(self.scale, rhs.scale),
            s: sub(self.s, rhs.s),
            m: sub(self.m, rhs.m),
            kg: sub(self.kg, rhs.kg),
            A: sub(self.A, rhs.A),
            K: sub(self.K, rhs.K),
            mol: sub(self.mol, rhs.mol),
            cd: sub(self.cd, rhs.cd),
        }
    }

    /// Raise SI units to an integer power
    ///
    /// ```
    /// # use const_units::units::{second, hertz};
    /// assert_eq!(second.powi(-1), hertz);
    /// ```
    pub const fn powi(self, v: i32) -> Self {
        self.powf((v, 1))
    }

    /// Raise SI units to a fractional power
    ///
    /// ```
    /// # use const_units::units::{meter};
    /// assert_eq!(meter.powi(2).powf((1, 2)), meter);
    /// ```
    ///
    /// # Panics
    /// Panics if `scale` can't be raised to the given power without losing precision
    ///
    /// ```should_panic
    /// # use const_units::units::{meter};
    /// meter.scale_by((2, 1)).powf((1, 2)); // `2` isn't a square number
    /// ```
    pub const fn powf(self, v: (i32, u32)) -> Self {
        match self.checked_powf(v) {
            Some(v) => v,
            None => panic!("Scale can't be raised to the given power without losing precision"),
        }
    }

    /// A checked version of `pow`
    ///
    /// ```
    /// # use const_units::units::{meter};
    /// assert_eq!(meter.powi(2).checked_powf((1, 2)), Some(meter));
    /// assert_eq!(meter.scale_by((2, 1)).checked_powf((1, 2)), None);
    /// ```
    pub const fn checked_powf(self, mut v: (i32, u32)) -> Option<Self> {
        v = simplify(v);
        let mut scale = self.scale;

        // `Try` isn't allowed in const
        scale.0 = match try_root(scale.0.abs() as u128, v.1) {
            Some(v) => v as i128 * scale.0.signum(),
            None => return None,
        };
        scale.1 = match try_root(scale.1, v.1) {
            Some(v) => v,
            None => return None,
        };

        if v.0.signum() == -1 {
            scale = recip(scale);
            v.0 = v.0.abs();
        }

        scale.0 = scale.0.pow(v.0 as u32);
        scale.1 = scale.1.pow(v.0 as u32);

        Some(SI {
            scale,
            s: simplify((self.s.0 * v.0, self.s.1 * v.1)),
            m: simplify((self.m.0 * v.0, self.m.1 * v.1)),
            kg: simplify((self.kg.0 * v.0, self.kg.1 * v.1)),
            A: simplify((self.A.0 * v.0, self.A.1 * v.1)),
            K: simplify((self.K.0 * v.0, self.K.1 * v.1)),
            mol: simplify((self.mol.0 * v.0, self.mol.1 * v.1)),
            cd: simplify((self.cd.0 * v.0, self.cd.1 * v.1)),
        })
    }
}

impl ConstParamTy for SI {}

/// A value with dimensionality
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Quantity<V, const UNITS: SI>(pub V);

impl<V: Display, const UNITS: SI> Display for Quantity<V, UNITS> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        self.0.fmt(f)?;
        f.write_str(" ")?;
        UNITS.fmt(f)?;

        Ok(())
    }
}

impl<V, const UNITS: SI> Quantity<V, UNITS> {
    /// Take the power to an integer
    ///
    /// ```
    /// # use const_units::{Quantity, units::{meter}};
    /// assert_eq!(Quantity::<f64, { meter }>(2.).powi::<2>(), Quantity::<f64, { meter.powi(2) }>(4.));
    /// ```
    pub fn powi<const EXP: i32>(self) -> Quantity<V::Output, { UNITS.powi(EXP) }>
    where
        V: Pow<i32>,
    {
        Quantity(self.0.pow(EXP))
    }

    /// Take the power to a fraction
    ///
    /// ```
    /// # use const_units::{Quantity, units::{meter}};
    /// assert_eq!(Quantity::<f64, { meter }>(4.).powf::<{(1, 2)}>(), Quantity::<f64, { meter.powf((1, 2)) }>(2.));
    /// ```
    pub fn powf<const EXP: (i32, u32)>(self) -> Quantity<V::Output, { UNITS.powf(EXP) }>
    where
        V: Pow<f64>,
    {
        Quantity(self.0.pow(EXP.0 as f64 / EXP.1 as f64))
    }
}

/// Create a conversion factor
///
/// Used to cancel out the scale factor without changing the underlying value
///
/// ```
/// # #![feature(generic_const_exprs)]
/// # use const_units::{Quantity, units::{DIMENSIONLESS, minute, second}, conversion_factor};
/// // Returns a `Quantity`
/// assert_eq!(conversion_factor::<f64, { (1, 60) }>(), Quantity::<f64, { DIMENSIONLESS.scale_by((1, 60)) }>(60.));
///
/// // Converting minutes to seconds
/// assert_eq!(Quantity::<f64, { minute }>(1.) * conversion_factor::<f64, { (1, 60) }>(), Quantity::<f64, { second }>(60.))
/// ```
pub fn conversion_factor<V, const SCALE: (i128, u128)>(
) -> Quantity<V, { DIMENSIONLESS.scale_by(SCALE) }>
where
    V: Div<V, Output = V> + NumCast,
{
    Quantity(
        V::from(SCALE.1).expect("Casting the scale value to type V to work")
            / V::from(SCALE.0).expect("Casting the scale value to type V to work"),
    )
}

/// Arbitrary math operations are only reasonable on dimensionless values
impl<V> Deref for Quantity<V, DIMENSIONLESS> {
    type Target = V;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Arbitrary math operations are only reasonable on dimensionless values
impl<V> DerefMut for Quantity<V, DIMENSIONLESS> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

/// Units must be the same for addition
impl<R, O, T: Add<R, Output = O>, const UNITS: SI> Add<Quantity<R, UNITS>> for Quantity<T, UNITS> {
    type Output = Quantity<O, UNITS>;

    fn add(self, rhs: Quantity<R, UNITS>) -> Self::Output {
        Quantity(self.0 + rhs.0)
    }
}

/// Units must be the same for addition
impl<R, O, T: Sub<R, Output = O>, const UNITS: SI> Sub<Quantity<R, UNITS>> for Quantity<T, UNITS> {
    type Output = Quantity<O, UNITS>;

    fn sub(self, rhs: Quantity<R, UNITS>) -> Self::Output {
        Quantity(self.0 - rhs.0)
    }
}

/// Units must be the same for addition
impl<R, T: AddAssign<R>, const UNITS: SI> AddAssign<Quantity<R, UNITS>> for Quantity<T, UNITS> {
    fn add_assign(&mut self, rhs: Quantity<R, UNITS>) {
        self.0 += rhs.0;
    }
}

/// Units must be the same for addition
impl<R, T: SubAssign<R>, const UNITS: SI> SubAssign<Quantity<R, UNITS>> for Quantity<T, UNITS> {
    fn sub_assign(&mut self, rhs: Quantity<R, UNITS>) {
        self.0 -= rhs.0;
    }
}

/// Scaling by a dimensionless value is OK
impl<R, T: MulAssign<R>, const UNITS: SI> MulAssign<Quantity<R, DIMENSIONLESS>>
    for Quantity<T, UNITS>
{
    fn mul_assign(&mut self, rhs: Quantity<R, DIMENSIONLESS>) {
        self.0 *= rhs.0;
    }
}

/// Scaling by a dimensionless value is OK
impl<R, T: DivAssign<R>, const UNITS: SI> DivAssign<Quantity<R, DIMENSIONLESS>>
    for Quantity<T, UNITS>
{
    fn div_assign(&mut self, rhs: Quantity<R, DIMENSIONLESS>) {
        self.0 /= rhs.0;
    }
}

// https://stackoverflow.com/questions/66361365/unconstrained-generic-constant-when-adding-const-generics

/// Multiplications multiplies the units
impl<R, O, T: Mul<R, Output = O>, const LHS_UNITS: SI, const RHS_UNITS: SI>
    Mul<Quantity<R, RHS_UNITS>> for Quantity<T, LHS_UNITS>
where
    Quantity<O, { SI::mul(LHS_UNITS, RHS_UNITS) }>: Sized,
{
    type Output = Quantity<O, { SI::mul(LHS_UNITS, RHS_UNITS) }>;

    fn mul(self, rhs: Quantity<R, RHS_UNITS>) -> Self::Output {
        Quantity(self.0 * rhs.0)
    }
}

/// Division divides the units
impl<R, O, T: Div<R, Output = O>, const LHS_UNITS: SI, const RHS_UNITS: SI>
    Div<Quantity<R, RHS_UNITS>> for Quantity<T, LHS_UNITS>
where
    Quantity<O, { SI::div(LHS_UNITS, RHS_UNITS) }>: Sized,
{
    type Output = Quantity<O, { SI::div(LHS_UNITS, RHS_UNITS) }>;

    fn div(self, rhs: Quantity<R, RHS_UNITS>) -> Self::Output {
        Quantity(self.0 / rhs.0)
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        add, gcd, mul, simplify, sub, try_root,
        units::{joule, meter, minute, mole, DIMENSIONLESS},
    };

    #[test]
    fn fractional_add() {
        assert_eq!(add((1, 2), (1, 3)), (5, 6));
        assert_eq!(add((1, 2), (-1, 3)), (1, 6));
        assert_eq!(add((-1, 2), (-1, 3)), (-5, 6));
        assert_eq!(add((-2, 2), (-1, 3)), (-4, 3));
        assert_eq!(add((2, 2), (3, 3)), (2, 1));
    }

    #[test]
    fn fractional_sub() {
        assert_eq!(sub((1, 2), (1, 3)), (1, 6));
        assert_eq!(sub((1, 2), (-1, 3)), (5, 6));
        assert_eq!(sub((-1, 2), (-1, 3)), (-1, 6));
        assert_eq!(sub((-2, 2), (-1, 3)), (-2, 3));
        assert_eq!(sub((2, 2), (3, 3)), (0, 1));
    }

    #[test]
    fn fractional_mul() {
        assert_eq!(mul((1, 2), (1, 3)), (1, 6));
        assert_eq!(mul((1, 2), (-1, 3)), (-1, 6));
        assert_eq!(mul((-1, 2), (-1, 3)), (1, 6));
        assert_eq!(mul((-2, 2), (-1, 3)), (1, 3));
        assert_eq!(mul((2, 2), (3, 3)), (1, 1));
        assert_eq!(mul((60, 1), (1, 60)), (1, 1));
    }

    #[test]
    fn test_gcd() {
        assert_eq!(gcd(2, 2), 2);
        assert_eq!(gcd(2, 3), 1);
        assert_eq!(gcd(2, 4), 2);
        assert_eq!(gcd(4, 3), 1);
        assert_eq!(gcd(4, 2), 2);
    }

    #[test]
    fn test_simplify() {
        assert_eq!(simplify((5, 6)), (5, 6));
        assert_eq!(simplify((-5, 6)), (-5, 6));
        assert_eq!(simplify((-4, 6)), (-2, 3));
        assert_eq!(simplify((4, 6)), (2, 3));
    }

    #[test]
    fn test_root() {
        assert_eq!(try_root(9, 2), Some(3));
        assert_eq!(try_root(27, 3), Some(3));
        assert_eq!(try_root(28, 3), None);
        assert_eq!(try_root(28, 39), None);
        assert_eq!(try_root(1, 39), Some(1));
        assert_eq!(try_root(28, 0), Some(1));
        assert_eq!(try_root(0, 1), Some(0));
    }

    #[test]
    fn test_display() {
        assert_eq!(format!("{}", mole), "mol");
        assert_eq!(format!("{}", joule), "m^2⋅kg⋅s^-2");
        assert_eq!(format!("{}", joule.powf((1, 2))), "m⋅kg^1/2⋅s^-1");
        assert_eq!(format!("{}", DIMENSIONLESS), "(1)");
        assert_eq!(format!("{}", DIMENSIONLESS.scale_by((1, 100))), "(1/100)");
        assert_eq!(format!("{}", meter.div(minute)), "(1/60)⋅m⋅s^-1");
    }
}