//! Implementation of modular arithmetic algorithms for all integer types in an
//! overflow-safe and const-compatible manner.

#![no_std]
#![warn(missing_docs)]

use core::ops::{Add, BitAnd, Index, IndexMut, Neg, Rem, Shr, Sub};
use num_traits::{one, CheckedAdd, One};

/// Represents an integer with an associated modulus.
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct Modular<T: Clone + Eq + Rem<Output = T>> {
    value: T,
    modulus: T,
}

impl<T: Clone + Eq + Rem<Output = T>> Modular<T> {
    /// Constructs a new modular arithmetic integer.
    pub fn new(value: T, modulus: T) -> Self {
        Self {
            value: value % modulus.clone(),
            modulus,
        }
    }

    /// Gets the value of the integer.
    pub fn value(&self) -> &T {
        &self.value
    }

    /// Gets the modulus of the integer.
    pub fn modulus(&self) -> &T {
        &self.modulus
    }

    /// Splits the integer into its value and modulus.
    pub fn into_parts(self) -> (T, T) {
        (self.value, self.modulus)
    }
}

impl<T: CheckedAdd + Clone + Eq + Rem<Output = T> + Shr<Output = T> + BitAnd<Output = T> + One>
    Modular<T>
{
    fn try_add(self, rhs: Self) -> Option<Self> {
        if self.modulus != rhs.modulus {
            None
        } else {
            Some(match self.value.checked_add(&rhs.value) {
                Some(r) => Self {
                    value: r % self.modulus.clone(),
                    modulus: self.modulus,
                },
                None => Self {
                    value: (((self.value.clone() >> one()) + (rhs.value.clone() >> one()))
                        % self.modulus.clone()
                        + ((self.value & one()) + (rhs.value & one())) % self.modulus.clone())
                        % self.modulus.clone(),
                    modulus: self.modulus,
                },
            })
        }
    }

    fn try_sub(self, rhs: Self) -> Option<Self>
    where
        T: Sub<Output = T>,
    {
        self.try_add(-rhs)
    }
}

impl<T: Clone + Eq + Rem<Output = T> + Sub<Output = T> + One> Neg for Modular<T> {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self {
            value: self.modulus.clone() - self.value - one(),
            modulus: self.modulus,
        }
    }
}

impl<T: CheckedAdd + Clone + Eq + Rem<Output = T> + Shr<Output = T> + BitAnd<Output = T> + One> Add
    for Modular<T>
{
    type Output = Modular<<T as Add>::Output>;

    fn add(self, rhs: Self) -> Self::Output {
        self.try_add(rhs)
            .expect("operands do not have the same modulus!")
    }
}

impl<
        T: CheckedAdd
            + Clone
            + Eq
            + Rem<Output = T>
            + Sub<Output = T>
            + Shr<Output = T>
            + BitAnd<Output = T>
            + One,
    > Sub for Modular<T>
{
    type Output = Modular<<T as Add>::Output>;

    fn sub(self, rhs: Self) -> Self::Output {
        self.try_sub(rhs)
            .expect("operands do not have the same modulus!")
    }
}

impl<T: Clone + Eq + Rem<Output = T> + One, U> Index<Modular<T>> for [U]
where
    [U]: Index<T, Output = U>,
{
    type Output = U;

    fn index(&self, index: Modular<T>) -> &Self::Output {
        &self[index.into_parts().0]
    }
}

impl<T: Clone + Eq + Rem<Output = T> + One, U> IndexMut<Modular<T>> for [U]
where
    [U]: IndexMut<T, Output = U>,
{
    fn index_mut(&mut self, index: Modular<T>) -> &mut Self::Output {
        &mut self[index.into_parts().0]
    }
}

/// Represents an integer with a statically specified modulus.
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct StaticModular<T: Clone + Eq + Rem<Output = T> + One + From<usize>, const M: usize>(T);

impl<T: Clone + Eq + Rem<Output = T> + One + From<usize>, const M: usize> StaticModular<T, M> {
    /// Constructs a new modular integer.
    pub fn new(value: T) -> Self {
        Self(value % M.into())
    }

    /// Gets the value of the integer.
    pub fn value(&self) -> &T {
        &self.0
    }

    /// Consumes the integer, producing its value.
    pub fn into_value(self) -> T {
        self.0
    }

    /// Consumes the integer, producing an equivalent [`Modular`] value.
    pub fn into_dynamic(self) -> Modular<T> {
        Modular {
            value: self.0,
            modulus: M.into(),
        }
    }
}

impl<
        T: CheckedAdd
            + Clone
            + Eq
            + Rem<Output = T>
            + Shr<Output = T>
            + BitAnd<Output = T>
            + One
            + From<usize>,
        const M: usize,
    > Add for StaticModular<T, M>
{
    type Output = StaticModular<<T as Add>::Output, M>;

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

impl<T: Clone + Eq + From<usize> + Rem<Output = T> + One, const M: usize> From<T>
    for StaticModular<T, M>
{
    fn from(value: T) -> Self {
        Self::new(value)
    }
}

impl<T: Clone + Eq + From<usize> + Rem<Output = T> + One, U, const M: usize>
    Index<StaticModular<T, M>> for [U]
where
    [U]: Index<T, Output = U>,
{
    type Output = U;

    fn index(&self, index: StaticModular<T, M>) -> &Self::Output {
        &self[index.into_value()]
    }
}

impl<T: Clone + Eq + From<usize> + Rem<Output = T> + One, U, const M: usize>
    IndexMut<StaticModular<T, M>> for [U]
where
    [U]: IndexMut<T, Output = U>,
{
    fn index_mut(&mut self, index: StaticModular<T, M>) -> &mut Self::Output {
        &mut self[index.into_value()]
    }
}