//! Variable-time fused square-and-reduce for sparse primes.
use super::U256;

impl U256 {
    /// Computes `self^2 mod modulus` for sparse primes in a single fused pass:
    /// schoolbook squaring directly into sparse reduce.
    ///
    /// Instead of materializing a [`U512`] and then reducing, this function
    /// keeps the 8-limb squared product in local variables and immediately
    /// applies the folding reduction `lo + hi * comp`. The squaring phase uses
    /// the symmetry optimization: 4 diagonal products plus 6 unique cross
    /// products (doubled), for 10 multiplications total instead of the 16
    /// needed by full schoolbook multiplication.
    ///
    /// `comp` is the sparse prime complement: `2^256 - modulus`, which must
    /// fit in a single `u64`.
    ///
    /// # Examples
    ///
    /// ```
    /// use cnfy_uint::u256::U256;
    /// let P = U256::from_be_limbs([
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    /// ]);
    ///
    /// let a = U256::from_be_limbs([0, 0, 0, 7]);
    /// assert_eq!(a.square_mod_sparse_vt(&P, 0x1000003D1), U256::from_be_limbs([0, 0, 0, 49]));
    /// ```
    #[inline]
    pub fn square_mod_sparse_vt(&self, modulus: &U256, comp: u64) -> U256 {
        let a = self.0;

        // === Phase 1: Squaring (10 products via symmetry) ===

        // Step 1: Accumulate unique cross products a[i]*a[j] for i < j.
        // Row 0: a[0] * a[1..3]
        let p = (a[0] as u128) * (a[1] as u128);
        let c1 = p as u64;
        let carry = (p >> 64) as u64;

        let p = (a[0] as u128) * (a[2] as u128) + (carry as u128);
        let c2 = p as u64;
        let carry = (p >> 64) as u64;

        let p = (a[0] as u128) * (a[3] as u128) + (carry as u128);
        let c3 = p as u64;
        let c4 = (p >> 64) as u64;

        // Row 1: a[1] * a[2..3], accumulate into c[3..5]
        let p = (a[1] as u128) * (a[2] as u128) + (c3 as u128);
        let c3 = p as u64;
        let carry = (p >> 64) as u64;

        let p = (a[1] as u128) * (a[3] as u128) + (c4 as u128) + (carry as u128);
        let c4 = p as u64;
        let c5 = (p >> 64) as u64;

        // Row 2: a[2] * a[3], accumulate into c[5..6]
        let p = (a[2] as u128) * (a[3] as u128) + (c5 as u128);
        let c5 = p as u64;
        let c6 = (p >> 64) as u64;

        // Step 2: Double all cross products (left shift by 1 bit).
        // Cross array is [0, c1, c2, c3, c4, c5, c6, 0] before doubling.
        // Process high-to-low so each `let` rebinding reads predecessors' old values.
        let c7 = c6 >> 63;
        let c6 = (c6 << 1) | (c5 >> 63);
        let c5 = (c5 << 1) | (c4 >> 63);
        let c4 = (c4 << 1) | (c3 >> 63);
        let c3 = (c3 << 1) | (c2 >> 63);
        let c2 = (c2 << 1) | (c1 >> 63);
        let c1 = c1 << 1;
        // c0 = 0 (no cross product at position 0)

        // Step 3: Add diagonal products a[i]^2 at positions 2*i, 2*i+1.
        // Diagonal a[0]^2 at positions [0, 1]
        let sq = (a[0] as u128) * (a[0] as u128);
        let r0 = sq as u64; // cross[0] = 0, so just sq_lo
        let sum = (c1 as u128) + (sq >> 64);
        let r1 = sum as u64;
        let carry = (sum >> 64) as u64;

        // Diagonal a[1]^2 at positions [2, 3]
        let sq = (a[1] as u128) * (a[1] as u128);
        let sum = (c2 as u128) + (sq as u64 as u128) + (carry as u128);
        let r2 = sum as u64;
        let sum = (c3 as u128) + (sq >> 64) + (sum >> 64);
        let r3 = sum as u64;
        let carry = (sum >> 64) as u64;

        // Diagonal a[2]^2 at positions [4, 5]
        let sq = (a[2] as u128) * (a[2] as u128);
        let sum = (c4 as u128) + (sq as u64 as u128) + (carry as u128);
        let r4 = sum as u64;
        let sum = (c5 as u128) + (sq >> 64) + (sum >> 64);
        let r5 = sum as u64;
        let carry = (sum >> 64) as u64;

        // Diagonal a[3]^2 at positions [6, 7]
        let sq = (a[3] as u128) * (a[3] as u128);
        let sum = (c6 as u128) + (sq as u64 as u128) + (carry as u128);
        let r6 = sum as u64;
        let sum = (c7 as u128) + (sq >> 64) + (sum >> 64);
        let r7 = sum as u64;

        // === Phase 2: Sparse reduce (inline) ===
        // lo = [r0, r1, r2, r3], hi = [r4, r5, r6, r7]
        // result ≡ lo + hi * comp (mod modulus)

        let c = comp as u128;

        // Round 1: fused multiply-accumulate of hi * comp + lo
        let sum = (r0 as u128) + c * (r4 as u128);
        let r0 = sum as u64;
        let carry = sum >> 64;

        let sum = (r1 as u128) + c * (r5 as u128) + carry;
        let r1 = sum as u64;
        let carry = sum >> 64;

        let sum = (r2 as u128) + c * (r6 as u128) + carry;
        let r2 = sum as u64;
        let carry = sum >> 64;

        let sum = (r3 as u128) + c * (r7 as u128) + carry;
        let r3 = sum as u64;
        let overflow = (sum >> 64) as u64;

        // Round 2: fold remaining overflow (typically ~33 bits for secp256k1)
        if overflow == 0 {
            return U256([r0, r1, r2, r3]).modulo(modulus);
        }

        let correction = (overflow as u128) * c;

        let sum = (r0 as u128) + (correction as u64 as u128);
        let r0 = sum as u64;
        let carry = sum >> 64;

        let sum = (r1 as u128) + (correction >> 64) + carry;
        let r1 = sum as u64;
        let carry = sum >> 64;

        let sum = (r2 as u128) + carry;
        let r2 = sum as u64;
        let carry = sum >> 64;

        let sum = (r3 as u128) + carry;
        let r3 = sum as u64;
        let overflow2 = (sum >> 64) as u64;

        let mut result = U256([r0, r1, r2, r3]);

        if (overflow2 != 0) | (result >= *modulus) {
            result = result.overflowing_sub(modulus).0;
        }

        result.modulo(modulus)
    }
}

#[cfg(test)]
mod ai_tests {
    use super::*;
    const P: U256 = U256::from_be_limbs([
        0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFEFFFFFC2F,
    ]);
    const COMP: u64 = 0x1000003D1;

    /// 7^2 = 49.
    #[test]
    fn small_square() {
        let a = U256::from_be_limbs([0, 0, 0, 7]);
        assert_eq!(
            a.square_mod_sparse_vt(&P, COMP),
            U256::from_be_limbs([0, 0, 0, 49])
        );
    }

    /// 0^2 mod P = 0.
    #[test]
    fn square_zero() {
        assert_eq!(U256::ZERO.square_mod_sparse_vt(&P, COMP), U256::ZERO);
    }

    /// 1^2 mod P = 1.
    #[test]
    fn square_one() {
        assert_eq!(U256::ONE.square_mod_sparse_vt(&P, COMP), U256::ONE);
    }

    /// (P-1)^2 mod P = 1.
    #[test]
    fn p_minus_one_squared() {
        let p_minus_1 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2E,
        ]);
        assert_eq!(p_minus_1.square_mod_sparse_vt(&P, COMP), U256::ONE);
    }

    /// Matches square_mod for diverse inputs.
    #[test]
    fn matches_square_mod() {
        let cases = [
            U256::from_be_limbs([0, 0, 0, 42]),
            U256::from_be_limbs([0, 0, 0x12345678, 0x9ABCDEF0]),
            U256::from_be_limbs([0x1111, 0x2222, 0x3333, 0x4444]),
            U256::from_be_limbs([
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x1111222233334444,
                0x5555666677778888,
            ]),
            U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]),
            U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, 0xFFFFFFFEFFFFFC2E]),
        ];
        for a in &cases {
            let expected = a.mul_mod(a, &P);
            let fused = a.square_mod_sparse_vt(&P, COMP);
            assert_eq!(fused, expected, "mismatch for {a:?}");
        }
    }

    /// Commutativity with mul_mod_sparse.
    #[test]
    fn matches_mul_mod_sparse() {
        let a = U256::from_be_limbs([
            0xAAAABBBBCCCCDDDD,
            0xEEEEFFFF00001111,
            0x2222333344445555,
            0x6666777788889999,
        ]);
        assert_eq!(
            a.square_mod_sparse_vt(&P, COMP),
            a.mul_mod_sparse(&a, &P, COMP)
        );
    }
}
