// TODO: make Zmod argument type generic
// TODO: re-implement ECM?

use std::fmt;
use std::ops;

// We assume canonical representative, can compare value for PartialEq
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct Zmod<const N: i64> {
    value: i64
}

fn canonical_rep<const N: i64>(x: i64) -> i64 {
    return (x % N + N) % N;
}

fn extended_gcd(a: i64, b: i64) -> (i64, i64, i64) {
    if b == 0 {
        return (a, 1, 0);
    }
    let (g, x, y) = extended_gcd(b, a%b);
    (g, y, x - y*(a/b))
}

impl<const N: i64> Zmod<N> {
    pub fn from(x: i64) -> Zmod<N> {
        #[cfg(debug_assertions)]
        assert!(N > 1, "modulus must be greater than 1");

        Zmod::<N> { value: canonical_rep::<N>(x) }
    }

    fn inverse(self) -> Result<Zmod<N>, i64> {
        let (g, a, _) = extended_gcd(self.value, N);
        if g == 1 {
            Ok(Zmod::<N>::from(a))
        } else {
            Err(g)
        }
    }
}

impl<const N: i64> fmt::Display for Zmod<N> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({} mod {})", self.value, N)
    }
}

impl<const N: i64> ops::Add for Zmod<N> {
    type Output = Zmod<N>;

    fn add(self, z: Zmod<N>) -> Zmod<N> {
        Zmod::<N>::from(self.value + z.value)
    }
}

impl<const N: i64> ops::Add<i64> for Zmod<N> {
    type Output = Zmod<N>;

    fn add(self, z: i64) -> Zmod<N> {
        Zmod::<N>::from(self.value + z)
    }
}

impl<const N: i64> ops::Add<Zmod<N>> for i64 {
    type Output = Zmod::<N>;

    fn add(self, z: Zmod::<N>) -> Zmod<N> {
        Zmod::<N>::from(self + z.value)
    }
}

impl<const N: i64> ops::AddAssign for Zmod<N> {
    fn add_assign(&mut self, z: Zmod<N>) {
        self.value = canonical_rep::<N>(self.value + z.value);
    }
}

impl<const N: i64> ops::AddAssign<i64> for Zmod<N> {
    fn add_assign(&mut self, z: i64) {
        self.value = canonical_rep::<N>(self.value + z);
    }
}

impl<const N: i64> ops::Sub for Zmod<N> {
    type Output = Zmod<N>;

    fn sub(self, z: Zmod<N>) -> Zmod<N> {
        Zmod::<N>::from(self.value - z.value)
    }
}

impl<const N: i64> ops::Sub<i64> for Zmod<N> {
    type Output = Zmod<N>;

    fn sub(self, z: i64) -> Zmod<N> {
        Zmod::<N>::from(self.value - z)
    }
}

impl<const N: i64> ops::Sub<Zmod<N>> for i64 {
    type Output = Zmod<N>;

    fn sub(self, z: Zmod<N>) -> Zmod<N> {
        Zmod::<N>::from(self - z.value)
    }
}

impl<const N: i64> ops::SubAssign for Zmod<N> {
    fn sub_assign(&mut self, z: Zmod<N>) {
        self.value = canonical_rep::<N>(self.value - z.value);
    }
}

impl<const N: i64> ops::SubAssign<i64> for Zmod<N> {
    fn sub_assign(&mut self, z: i64) {
        self.value = canonical_rep::<N>(self.value - z);
    }
}

impl<const N: i64> ops::Neg for Zmod<N> {
    type Output = Zmod<N>;

    fn neg(self) -> Zmod<N> {
        Zmod::<N>::from(-self.value)
    }
}

impl<const N: i64> ops::Mul for Zmod<N> {
    type Output = Zmod<N>;

    fn mul(self, z: Zmod<N>) -> Zmod<N> {
        Zmod::<N>::from(self.value * z.value)
    }
}

impl<const N: i64> ops::Mul<i64> for Zmod<N> {
    type Output = Zmod<N>;

    fn mul(self, z: i64) -> Zmod<N> {
        Zmod::<N>::from(self.value * z)
    }
}

impl<const N: i64> ops::Mul<Zmod<N>> for i64 {
    type Output = Zmod<N>;

    fn mul(self, z: Zmod<N>) -> Zmod<N> {
        Zmod::<N>::from(self * z.value)
    }
}

impl<const N: i64> ops::MulAssign for Zmod<N> {
    fn mul_assign(&mut self, z: Zmod<N>) {
        self.value = canonical_rep::<N>(self.value * z.value);
    }
}

impl<const N: i64> ops::MulAssign<i64> for Zmod<N> {
    fn mul_assign(&mut self, z: i64) {
        self.value = canonical_rep::<N>(self.value * z);
    }
}

impl<const N: i64> ops::Div for Zmod<N> {
    type Output = Result<Zmod<N>, i64>;

    fn div(self, z: Zmod<N>) -> Result<Zmod<N>, i64> {
        Ok(self * z.inverse()?)
    }
}

impl<const N: i64> ops::Div<i64> for Zmod<N> {
    type Output = Result<Zmod<N>, i64>;

    fn div(self, z: i64) -> Result<Zmod<N>, i64> {
        self / Zmod::<N>::from(z)
    }
}

impl<const N: i64> ops::Div<Zmod<N>> for i64 {
    type Output = Result<Zmod<N>, i64>;

    fn div(self, z: Zmod<N>) -> Result<Zmod<N>, i64> {
        Zmod::<N>::from(self) / z
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fmt_simple() {
        let x = Zmod::<5>::from(3);
        assert_eq!(x.to_string(), "(3 mod 5)");
    }

    #[test]
    fn two_is_zero_mod_two() {
        assert_eq!(Zmod::<2>::from(2), Zmod::<2>::from(0));
    }

    #[test]
    fn negative_one_is_one_mod_two() {
        assert_eq!(Zmod::<2>::from(-1), Zmod::<2>::from(1));
    }

    #[test]
    #[should_panic]
    fn negative_modulus_panic() {
        let _ = Zmod::<-3>::from(0);
    }

    #[test]
    #[should_panic]
    fn modulus_one_panic() {
        let _ = Zmod::<1>::from(0);
    }

    #[test]
    fn add_zmod_zmod() {
        let one = Zmod::<2>::from(1);
        assert_eq!(one + one, Zmod::<2>::from(0));
    }

    #[test]
    fn add_zmod_num() {
        let one = Zmod::<4>::from(1);
        assert_eq!(one + 3, Zmod::<4>::from(0));
    }

    #[test]
    fn add_num_zmod() {
        assert_eq!(3 + Zmod::<9>::from(-4), Zmod::<9>::from(8));
    }

    #[test]
    fn add_assign_zmod_zmod() {
        let mut x = Zmod::<7>::from(-2);
        x += Zmod::<7>::from(25);
        assert_eq!(x, Zmod::<7>::from(2));
    }

    #[test]
    fn add_assign_zmod_num() {
        let mut x = Zmod::<3>::from(2);
        x += 2;
        assert_eq!(x, Zmod::<3>::from(1));
    }

    #[test]
    fn subtract_zmod_zmod() {
        let x = Zmod::<5>::from(2);
        let y = Zmod::<5>::from(-4);
        assert_eq!(x - y, Zmod::<5>::from(1));
    }

    #[test]
    fn subtract_zmod_num() {
        assert_eq!(Zmod::<3>::from(1) - 2, Zmod::<3>::from(-1));
    }

    #[test]
    fn subtract_num_zmod() {
        assert_eq!(2 - Zmod::<7>::from(5), Zmod::<7>::from(4));
    }

    #[test]
    fn subtract_assign_zmod_zmod() {
        let mut x = Zmod::<15>::from(32);
        x -= Zmod::<15>::from(12);
        assert_eq!(x, Zmod::<15>::from(20));
    }

    #[test]
    fn subtract_assign_zmod_num() {
        let mut x = Zmod::<17>::from(11);
        x -= 20;
        assert_eq!(x, Zmod::<17>::from(-9));
    }

    #[test]
    fn neg() {
        assert_eq!(-Zmod::<14>::from(18), Zmod::<14>::from(-18));
    }

    #[test]
    fn multiply_zmod_zmod() {
        let x = Zmod::<9>::from(6);
        let y = Zmod::<9>::from(-3);
        assert_eq!(x * y, Zmod::<9>::from(0));
    }

    #[test]
    fn multiply_zmod_num() {
        assert_eq!(Zmod::<5>::from(4) * 2, Zmod::<5>::from(3));
    }

    #[test]
    fn multiply_num_zmod() {
        assert_eq!(6 * Zmod::<7>::from(5), Zmod::<7>::from(30));
    }

    #[test]
    fn multiply_assign_zmod_zmod() {
        let mut x = Zmod::<15>::from(3);
        x *= Zmod::<15>::from(7);
        assert_eq!(x, Zmod::<15>::from(6));
    }

    #[test]
    fn multiply_assign_zmod_num() {
        let mut x = Zmod::<17>::from(11);
        x *= 4;
        assert_eq!(x, Zmod::<17>::from(10));
    }

    #[test]
    fn inverse_one() {
        let x = Zmod::<44>::from(1);
        assert_eq!(x.inverse(), Ok(x));
    }

    #[test]
    fn inverse_12_35() {
        let x = Zmod::<35>::from(12);
        assert_eq!(x.inverse(), Ok(Zmod::<35>::from(3)));
    }

    #[test]
    fn inverse_fail() {
        let x = Zmod::<9>::from(6);
        assert_eq!(x.inverse(), Err(3));
    }

    #[test]
    fn divide_success() {
        let x = Zmod::<35>::from(15);
        let d = Zmod::<35>::from(6);
        assert_eq!(x / d, Ok(Zmod::<35>::from(20)));
    }
}