//! Prime-field arithmetic for p < 2^62 in Montgomery form (R = 2^64).
//!
//! All field elements handled by the walks are kept in Montgomery representation
//! `a·R mod p`; conversion happens only at the boundaries (CLI input/output, JSON) and
//! where a canonical (standard-form) value is needed for an ordering decision.
//! Products of two residues fit in a u128; the Montgomery reduction needs one
//! conditional subtraction because p < 2^62 guarantees (a·b + m·p)/R < 2p.
//! Exact integer arithmetic only (CONVENTIONS.md §1).

#[derive(Clone, Copy, Debug)]
pub struct Field {
    pub p: u64,
    /// -p^{-1} mod 2^64
    pub pinv: u64,
    /// R mod p   (Montgomery form of 1)
    pub one: u64,
    /// R^2 mod p
    pub r2: u64,
    /// R^3 mod p (used to convert a plain inverse back to Montgomery form)
    pub r3: u64,
}

impl Field {
    pub fn new(p: u64) -> Field {
        assert!(p > 3 && p & 1 == 1, "p must be an odd prime > 3");
        assert!(p < (1u64 << 62), "p must be < 2^62");
        // Newton iteration for p^{-1} mod 2^64
        let mut inv: u64 = 1;
        for _ in 0..6 {
            inv = inv.wrapping_mul(2u64.wrapping_sub(p.wrapping_mul(inv)));
        }
        debug_assert_eq!(p.wrapping_mul(inv), 1);
        let pinv = inv.wrapping_neg();
        let r = ((1u128 << 64) % p as u128) as u64;
        let r2 = ((r as u128 * r as u128) % p as u128) as u64;
        let r3 = ((r2 as u128 * r as u128) % p as u128) as u64;
        Field { p, pinv, one: r, r2, r3 }
    }

    /// Montgomery reduction of t < p·2^64: returns t·R^{-1} mod p.
    #[inline(always)]
    pub fn redc(&self, t: u128) -> u64 {
        let m = (t as u64).wrapping_mul(self.pinv);
        let u = ((t + (m as u128) * (self.p as u128)) >> 64) as u64;
        if u >= self.p {
            u - self.p
        } else {
            u
        }
    }

    #[inline(always)]
    pub fn mul(&self, a: u64, b: u64) -> u64 {
        self.redc(a as u128 * b as u128)
    }

    #[inline(always)]
    pub fn sqr(&self, a: u64) -> u64 {
        self.redc(a as u128 * a as u128)
    }

    #[inline(always)]
    pub fn add(&self, a: u64, b: u64) -> u64 {
        let s = a + b; // < 2^63, no overflow
        if s >= self.p {
            s - self.p
        } else {
            s
        }
    }

    #[inline(always)]
    pub fn sub(&self, a: u64, b: u64) -> u64 {
        if a >= b {
            a - b
        } else {
            a + self.p - b
        }
    }

    #[inline(always)]
    pub fn neg(&self, a: u64) -> u64 {
        if a == 0 {
            0
        } else {
            self.p - a
        }
    }

    #[inline(always)]
    pub fn dbl(&self, a: u64) -> u64 {
        self.add(a, a)
    }

    /// Standard integer a (0 ≤ a < p) → Montgomery form.
    #[inline(always)]
    pub fn to_mont(&self, a: u64) -> u64 {
        debug_assert!(a < self.p);
        self.mul(a, self.r2)
    }

    /// Montgomery form → standard integer.
    #[inline(always)]
    pub fn from_mont(&self, a: u64) -> u64 {
        self.redc(a as u128)
    }

    /// Inverse of a standard-form element by the extended Euclidean algorithm (a ≠ 0).
    pub fn inv_plain(&self, a: u64) -> u64 {
        debug_assert!(a != 0 && a < self.p);
        let (mut r0, mut r1) = (self.p as i64, a as i64);
        let (mut t0, mut t1) = (0i64, 1i64);
        while r1 != 0 {
            let q = r0 / r1;
            let r2 = r0 - q * r1;
            r0 = r1;
            r1 = r2;
            let t2 = t0 - q * t1;
            t0 = t1;
            t1 = t2;
        }
        debug_assert_eq!(r0, 1, "inverse of non-unit");
        if t0 < 0 {
            (t0 + self.p as i64) as u64
        } else {
            t0 as u64
        }
    }

    /// Inverse of a Montgomery-form element, result in Montgomery form.
    /// inv_plain(aR) = a^{-1} R^{-1}; multiplying by R^3 (one Montgomery product) gives a^{-1} R.
    #[inline]
    pub fn inv(&self, a: u64) -> u64 {
        let t = self.inv_plain(a);
        self.mul(t, self.r3)
    }

    /// Montgomery's simultaneous inversion: replaces every element of `xs` (all non-zero,
    /// Montgomery form) by its inverse, using 3(k−1) multiplications and a single inversion.
    /// `scratch` must have length ≥ xs.len(). Returns (muls, inversions) performed.
    pub fn batch_inv(&self, xs: &mut [u64], scratch: &mut [u64]) -> (u64, u64) {
        let k = xs.len();
        if k == 0 {
            return (0, 0);
        }
        scratch[0] = xs[0];
        for i in 1..k {
            scratch[i] = self.mul(scratch[i - 1], xs[i]);
        }
        let mut acc = self.inv(scratch[k - 1]);
        for i in (1..k).rev() {
            let inv_i = self.mul(acc, scratch[i - 1]);
            acc = self.mul(acc, xs[i]);
            xs[i] = inv_i;
        }
        xs[0] = acc;
        ((3 * (k as u64 - 1)) + 1, 1)
    }

    pub fn pow_plain(&self, a: u64, mut e: u64) -> u64 {
        let mut base = self.to_mont(a);
        let mut acc = self.one;
        while e > 0 {
            if e & 1 == 1 {
                acc = self.mul(acc, base);
            }
            base = self.sqr(base);
            e >>= 1;
        }
        self.from_mont(acc)
    }

    /// Legendre symbol as 0, 1 or p−1 (standard form input).
    pub fn legendre(&self, a: u64) -> u64 {
        self.pow_plain(a % self.p, (self.p - 1) / 2)
    }

    /// Tonelli–Shanks square root of a standard-form element; None if non-residue.
    pub fn sqrt_plain(&self, a: u64) -> Option<u64> {
        let p = self.p;
        let a = a % p;
        if a == 0 {
            return Some(0);
        }
        if self.legendre(a) != 1 {
            return None;
        }
        if p % 4 == 3 {
            return Some(self.pow_plain(a, (p + 1) / 4));
        }
        let (mut q, mut s) = (p - 1, 0u32);
        while q % 2 == 0 {
            q /= 2;
            s += 1;
        }
        let mut z = 2u64;
        while self.legendre(z) != p - 1 {
            z += 1;
        }
        let mut m = s;
        let mut c = self.to_mont(self.pow_plain(z, q));
        let mut t = self.to_mont(self.pow_plain(a, q));
        let mut r = self.to_mont(self.pow_plain(a, (q + 1) / 2));
        while t != self.one {
            let mut i = 0u32;
            let mut t2 = t;
            while t2 != self.one {
                t2 = self.sqr(t2);
                i += 1;
            }
            let mut b = c;
            for _ in 0..(m - i - 1) {
                b = self.sqr(b);
            }
            m = i;
            c = self.sqr(b);
            t = self.mul(t, c);
            r = self.mul(r, b);
        }
        Some(self.from_mont(r))
    }
}

/// Arithmetic modulo the (prime) group order n < 2^62 — used for the scalar bookkeeping.
#[derive(Clone, Copy, Debug)]
pub struct Modn {
    pub n: u64,
}

impl Modn {
    pub fn new(n: u64) -> Modn {
        assert!(n > 1 && n < (1u64 << 62));
        Modn { n }
    }
    #[inline(always)]
    pub fn add(&self, a: u64, b: u64) -> u64 {
        let s = a + b;
        if s >= self.n {
            s - self.n
        } else {
            s
        }
    }
    #[inline(always)]
    pub fn sub(&self, a: u64, b: u64) -> u64 {
        if a >= b {
            a - b
        } else {
            a + self.n - b
        }
    }
    #[inline(always)]
    pub fn neg(&self, a: u64) -> u64 {
        if a == 0 {
            0
        } else {
            self.n - a
        }
    }
    #[inline(always)]
    pub fn mul(&self, a: u64, b: u64) -> u64 {
        ((a as u128 * b as u128) % self.n as u128) as u64
    }
    #[inline(always)]
    pub fn dbl(&self, a: u64) -> u64 {
        self.add(a, a)
    }
    pub fn inv(&self, a: u64) -> Option<u64> {
        if a == 0 {
            return None;
        }
        let (mut r0, mut r1) = (self.n as i64, a as i64);
        let (mut t0, mut t1) = (0i64, 1i64);
        while r1 != 0 {
            let q = r0 / r1;
            let r2 = r0 - q * r1;
            r0 = r1;
            r1 = r2;
            let t2 = t0 - q * t1;
            t0 = t1;
            t1 = t2;
        }
        if r0 != 1 {
            return None;
        }
        Some(if t0 < 0 { (t0 + self.n as i64) as u64 } else { t0 as u64 })
    }
    /// Solve a·d ≡ b (mod n); None if a ≡ 0.
    pub fn solve(&self, a: u64, b: u64) -> Option<u64> {
        self.inv(a).map(|ai| self.mul(ai, b))
    }
    pub fn pow(&self, a: u64, mut e: u64) -> u64 {
        let mut base = a % self.n;
        let mut acc = 1u64;
        while e > 0 {
            if e & 1 == 1 {
                acc = self.mul(acc, base);
            }
            base = self.mul(base, base);
            e >>= 1;
        }
        acc
    }
}

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

    #[test]
    fn montgomery_roundtrip_and_inverse() {
        for &p in &[1_000_003u64, 4611686018427387847u64 /* 2^62-57 */, 2305843009213693951u64] {
            let f = Field::new(p);
            let mut x = 123456789u64 % p;
            for _ in 0..1000 {
                x = (x as u128 * 6364136223846793005u128 % p as u128) as u64 + 1;
                x %= p;
                if x == 0 {
                    continue;
                }
                let xm = f.to_mont(x);
                assert_eq!(f.from_mont(xm), x);
                let y = f.inv(xm);
                assert_eq!(f.from_mont(f.mul(xm, y)), 1);
                let prod = f.from_mont(f.mul(xm, xm));
                assert_eq!(prod, ((x as u128 * x as u128) % p as u128) as u64);
            }
        }
    }

    #[test]
    fn batch_inverse_matches_single() {
        let p = 4611686018427387847u64;
        let f = Field::new(p);
        let mut xs: Vec<u64> = (1..300u64).map(|i| f.to_mont(i * 987654321 % p)).collect();
        let expect: Vec<u64> = xs.iter().map(|&x| f.inv(x)).collect();
        let mut scratch = vec![0u64; xs.len()];
        f.batch_inv(&mut xs, &mut scratch);
        assert_eq!(xs, expect);
    }

    #[test]
    fn sqrt_works() {
        for &p in &[1_000_003u64, 1_000_033u64, 4611686018427387847u64] {
            let f = Field::new(p);
            for a in 2..200u64 {
                if let Some(r) = f.sqrt_plain(a) {
                    assert_eq!(((r as u128 * r as u128) % p as u128) as u64, a % p);
                }
            }
        }
    }
}