//! Bézout coefficient `(d, e)` matrix application for [`DivstepMatrix`].
use super::DivstepMatrix;
use crate::u256::U256;

impl DivstepMatrix {
    /// Applies the accumulated transition matrix to Bézout coefficients `d`
    /// and `e` with modular reduction, then divides by `2^62` modulo `modulus`.
    ///
    /// Computes:
    /// ```text
    /// d' = (u*d + v*e) / 2^62  mod modulus
    /// e' = (q*d + r*e) / 2^62  mod modulus
    /// ```
    ///
    /// Division by `2^62 mod modulus` is performed by repeated halving:
    /// if the value is odd, add modulus to make it even, then shift right.
    /// The carry from the addition is incorporated into the most significant
    /// bit after the shift.
    pub(crate) fn apply_de(
        &self,
        d: &U256,
        e: &U256,
        modulus: &U256,
    ) -> (U256, U256) {
        let new_d = Self::signed_mod_linear_combo(self.u, self.v, d, e, modulus);
        let new_e = Self::signed_mod_linear_combo(self.q, self.r, d, e, modulus);

        // Divide by 2^62 via repeated halving mod modulus
        let new_d = Self::div_pow2_mod(new_d, 62, modulus);
        let new_e = Self::div_pow2_mod(new_e, 62, modulus);

        (new_d, new_e)
    }

    /// Computes `(a * x + b * y) mod modulus` where `a`, `b` are signed i64
    /// and `x`, `y` are U256 values in `[0, modulus)`.
    fn signed_mod_linear_combo(
        a: i64,
        b: i64,
        x: &U256,
        y: &U256,
        modulus: &U256,
    ) -> U256 {
        let ax = Self::signed_mul_mod(a, x, modulus);
        let by = Self::signed_mul_mod(b, y, modulus);
        ax.add_mod(&by, modulus)
    }

    /// Computes `scalar * value mod modulus` where `scalar` is a signed i64.
    ///
    /// For negative scalars, computes `(-scalar) * value` then negates the
    /// result modulo `modulus`.
    fn signed_mul_mod(scalar: i64, value: &U256, modulus: &U256) -> U256 {
        let (abs_scalar, neg) = if scalar < 0 {
            (scalar.unsigned_abs(), true)
        } else {
            (scalar as u64, false)
        };

        let product = Self::mul_u64_mod(abs_scalar, value, modulus);

        if neg {
            U256::ZERO.sub_mod(&product, modulus)
        } else {
            product
        }
    }

    /// Multiplies a U256 by a u64 modulo `modulus`.
    ///
    /// Decomposes the u64 into high and low 32-bit halves and uses
    /// mul_mod to stay reduced.
    fn mul_u64_mod(scalar: u64, value: &U256, modulus: &U256) -> U256 {
        let lo = scalar & 0xFFFFFFFF;
        let hi = scalar >> 32;

        let val_lo = U256::from_be_limbs([0, 0, 0, lo]).mul_mod(value, modulus);
        if hi == 0 {
            return val_lo;
        }
        let val_hi = U256::from_be_limbs([0, 0, 0, hi]).mul_mod(value, modulus);
        let shift = U256::from_be_limbs([0, 0, 0, 1u64 << 32]);
        let val_hi_shifted = val_hi.mul_mod(&shift, modulus);
        val_hi_shifted.add_mod(&val_lo, modulus)
    }

    /// Divides `value` by `2^n` modulo `modulus` via repeated conditional halving.
    ///
    /// For each of the `n` divisions by 2: if `value` is odd, adds `modulus`
    /// (making it even), then right-shifts by 1. The carry from the addition
    /// is incorporated as the new MSB after the shift, since `value + modulus`
    /// can exceed 256 bits.
    fn div_pow2_mod(mut value: U256, n: u32, modulus: &U256) -> U256 {
        for _ in 0..n {
            let is_odd = value.0[0] & 1;
            let (with_mod, carry) = value.overflowing_add(modulus);
            // Select: with_mod when odd, value when even
            value = with_mod.ct_select(&value, is_odd != 0);
            // Carry only applies when we chose the with_mod path
            let carry_bit = (carry as u64) & is_odd;
            // Right-shift by 1 with carry feeding into MSB
            value = U256([
                (value.0[0] >> 1) | (value.0[1] << 63),
                (value.0[1] >> 1) | (value.0[2] << 63),
                (value.0[2] >> 1) | (value.0[3] << 63),
                (value.0[3] >> 1) | (carry_bit << 63),
            ]);
        }
        value
    }
}

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

    const P: U256 = U256::from_be_limbs([
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    ]);

    /// div_pow2_mod(2, 1) = 1.
    #[test]
    fn div2_of_two_is_one() {
        let two = U256::from_be_limbs([0, 0, 0, 2]);
        assert_eq!(DivstepMatrix::div_pow2_mod(two, 1, &P), U256::ONE);
    }

    /// div_pow2_mod(1, 1) = (P+1)/2 (modular half of 1).
    #[test]
    fn div2_of_one() {
        let half = DivstepMatrix::div_pow2_mod(U256::ONE, 1, &P);
        let two = U256::from_be_limbs([0, 0, 0, 2]);
        assert_eq!(half.mul_mod(&two, &P), U256::ONE);
    }

    /// div_pow2_mod with value near P (exercises carry from overflow).
    #[test]
    fn div2_near_p_exercises_carry() {
        // P - 2 is even, so div2 should give (P-2)/2
        let p_minus_2 = P - U256::from_be_limbs([0, 0, 0, 2]);
        let half = DivstepMatrix::div_pow2_mod(p_minus_2, 1, &P);
        let two = U256::from_be_limbs([0, 0, 0, 2]);
        assert_eq!(half.mul_mod(&two, &P), p_minus_2);
    }

    /// div_pow2_mod of P-1 (odd, large) exercises carry path.
    #[test]
    fn div2_of_p_minus_one() {
        let p_minus_1 = P - U256::ONE;
        let half = DivstepMatrix::div_pow2_mod(p_minus_1, 1, &P);
        // 2 * half should equal P-1 mod P
        let two = U256::from_be_limbs([0, 0, 0, 2]);
        assert_eq!(half.mul_mod(&two, &P), p_minus_1);
    }

    /// signed_mul_mod with positive scalar.
    #[test]
    fn signed_mul_mod_positive() {
        let val = U256::from_be_limbs([0, 0, 0, 7]);
        let result = DivstepMatrix::signed_mul_mod(3, &val, &P);
        assert_eq!(result, U256::from_be_limbs([0, 0, 0, 21]));
    }

    /// signed_mul_mod with negative scalar.
    #[test]
    fn signed_mul_mod_negative() {
        let val = U256::from_be_limbs([0, 0, 0, 7]);
        let result = DivstepMatrix::signed_mul_mod(-1, &val, &P);
        let expected = P - U256::from_be_limbs([0, 0, 0, 7]);
        assert_eq!(result, expected);
    }

    /// Identity matrix apply_de leaves d, e unchanged (modulo 2^62 division).
    #[test]
    fn identity_preserves_values() {
        let m = DivstepMatrix::identity();
        let d = U256::from_be_limbs([0, 0, 0, 1u64 << 62]);
        let e = U256::from_be_limbs([0, 0, 0, 1u64 << 62]);
        let (new_d, new_e) = m.apply_de(&d, &e, &P);
        assert_eq!(new_d, U256::ONE);
        assert_eq!(new_e, U256::ONE);
    }

    /// div_pow2_mod roundtrip: (x / 2^62) * 2^62 = x mod P.
    #[test]
    fn div_pow2_roundtrip() {
        let val = U256::from_be_limbs([0, 0, 0, 42]);
        let divided = DivstepMatrix::div_pow2_mod(val, 62, &P);
        // Multiply back by 2^62
        let pow2_62 = U256::from_be_limbs([0, 0, 0, 1u64 << 62]);
        let recovered = divided.mul_mod(&pow2_62, &P);
        assert_eq!(recovered, val);
    }
}
