//! Polynomial operations over a prime field for secret sharing.

/// A large prime for field arithmetic (Mersenne-like prime).
/// Using 2^61 - 1 for efficient modular arithmetic.
pub const PRIME: u64 = (1u64 << 61) - 1;

/// Modular addition.
pub fn mod_add(a: u64, b: u64) -> u64 {
    let sum = a.wrapping_add(b);
    if sum >= PRIME || sum < a { sum.wrapping_sub(PRIME) } else { sum }
}

/// Modular subtraction.
pub fn mod_sub(a: u64, b: u64) -> u64 {
    if a >= b { a - b } else { PRIME - (b - a) }
}

/// Modular multiplication using 128-bit intermediate.
pub fn mod_mul(a: u64, b: u64) -> u64 {
    let wide = (a as u128) * (b as u128);
    mod_reduce(wide)
}

/// Reduce a 128-bit value modulo PRIME.
fn mod_reduce(wide: u128) -> u64 {
    // For PRIME = 2^61 - 1: x mod (2^61-1) = (x & PRIME) + (x >> 61)
    // May need two rounds
    let mut x = wide;
    loop {
        let lo = (x & (PRIME as u128)) as u64;
        let hi = (x >> 61) as u64;
        let sum = lo.wrapping_add(hi);
        if sum >= PRIME {
            x = (sum - PRIME) as u128;
        } else {
            return sum;
        }
    }
}

/// Modular exponentiation (square-and-multiply).
pub fn mod_pow(base: u64, exp: u64) -> u64 {
    if exp == 0 { return 1; }
    let mut result = 1u64;
    let mut b = base;
    let mut e = exp;
    while e > 0 {
        if e & 1 == 1 {
            result = mod_mul(result, b);
        }
        b = mod_mul(b, b);
        e >>= 1;
    }
    result
}

/// Modular inverse using Fermat's little theorem.
/// Since PRIME is prime, a^(-1) = a^(p-2) mod p.
pub fn mod_inv(a: u64) -> u64 {
    mod_pow(a, PRIME - 2)
}

/// Evaluate a polynomial at point x.
/// coeffs[0] is the constant term (the secret).
pub fn eval_poly(coeffs: &[u64], x: u64) -> u64 {
    if x == 0 {
        return coeffs[0];
    }
    // Horner's method: a0 + x*(a1 + x*(a2 + ...))
    let mut result = 0u64;
    for &coeff in coeffs.iter().rev() {
        result = mod_add(mod_mul(result, x), coeff);
    }
    result
}

/// Lagrange interpolation at x=0 to recover the constant term.
/// Given points (x_i, y_i), compute f(0).
pub fn lagrange_interpolate_zero(points: &[(u64, u64)]) -> u64 {
    let mut secret = 0u64;
    for (i, &(xi, yi)) in points.iter().enumerate() {
        let mut numerator = 1u64;
        let mut denominator = 1u64;
        for (j, &(xj, _)) in points.iter().enumerate() {
            if i != j {
                // For f(0): numerator *= (0 - xj) = PRIME - xj
                numerator = mod_mul(numerator, mod_sub(0, xj));
                denominator = mod_mul(denominator, mod_sub(xi, xj));
            }
        }
        let lagrange_coeff = mod_mul(numerator, mod_inv(denominator));
        secret = mod_add(secret, mod_mul(yi, lagrange_coeff));
    }
    secret
}

/// Lagrange interpolation at an arbitrary point x.
pub fn lagrange_interpolate(points: &[(u64, u64)], x: u64) -> u64 {
    let mut result = 0u64;
    for (i, &(xi, yi)) in points.iter().enumerate() {
        let mut basis = 1u64;
        for (j, &(xj, _)) in points.iter().enumerate() {
            if i != j {
                let num = mod_sub(x, xj);
                let den = mod_sub(xi, xj);
                basis = mod_mul(basis, mod_mul(num, mod_inv(den)));
            }
        }
        result = mod_add(result, mod_mul(yi, basis));
    }
    result
}

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

    #[test]
    fn test_mod_add() {
        let a = mod_add(PRIME - 1, 1);
        assert_eq!(a, 0);
        let b = mod_add(100, 200);
        assert_eq!(b, 300);
    }

    #[test]
    fn test_mod_mul_identity() {
        assert_eq!(mod_mul(1, 42), 42);
        assert_eq!(mod_mul(42, 1), 42);
    }

    #[test]
    fn test_mod_inv() {
        let a = 12345u64;
        let inv = mod_inv(a);
        assert_eq!(mod_mul(a, inv), 1);
    }

    #[test]
    fn test_mod_pow() {
        assert_eq!(mod_pow(2, 10), 1024);
        assert_eq!(mod_pow(3, 0), 1);
    }

    #[test]
    fn test_eval_poly_constant() {
        // f(x) = 42
        assert_eq!(eval_poly(&[42], 5), 42);
        assert_eq!(eval_poly(&[42], 0), 42);
    }

    #[test]
    fn test_eval_poly_linear() {
        // f(x) = 10 + 5*x (mod PRIME)
        assert_eq!(eval_poly(&[10, 5], 0), 10);
        assert_eq!(eval_poly(&[10, 5], 1), 15);
        assert_eq!(eval_poly(&[10, 5], 2), 20);
    }

    #[test]
    fn test_lagrange_constant() {
        let points = vec![(1u64, 42u64), (2, 42), (3, 42)];
        assert_eq!(lagrange_interpolate_zero(&points), 42);
    }

    #[test]
    fn test_lagrange_linear() {
        // f(x) = 10 + 5x
        let points = vec![(1u64, 15u64), (2u64, 20u64)];
        assert_eq!(lagrange_interpolate_zero(&points), 10);
    }
}
