//! Full 512-bit squaring using the schoolbook symmetry optimization.
use super::U256;
use crate::u512::U512;

impl U256 {
    /// Computes the full 512-bit square `self * self` using the schoolbook
    /// squaring optimization that exploits symmetry in the product matrix.
    ///
    /// For a 4-limb number, full multiplication requires 16 limb products.
    /// Squaring reduces this to 10: 4 diagonal products (`a[i] * a[i]`)
    /// plus 6 unique cross products (`a[i] * a[j]` for `i < j`), each
    /// doubled. The cross products are accumulated, shifted left by one
    /// bit, then the diagonal products are added with carry propagation.
    ///
    /// Returns the result as a [`U512`] in big-endian limb order.
    ///
    /// # Examples
    ///
    /// ```
    /// use cnfy_uint::u256::U256;
    /// use cnfy_uint::u512::U512;
    ///
    /// let a = U256::from_be_limbs([0, 0, 0, 7]);
    /// assert_eq!(a.square(), U512::from_be_limbs([0, 0, 0, 0, 0, 0, 0, 49]));
    /// ```
    #[inline]
    pub fn square(&self) -> U512 {
        let a = self.0;
        let mut cross = [0u64; 8];

        // Step 1: Accumulate cross products (i < j) with carry propagation.
        // Each product a[i]*a[j] fits in u128. Adding existing accumulator
        // (u64) plus carry (u64) still fits: (2^64-1)^2 + 2*(2^64-1) = 2^128-1.
        let mut carry: u64 = 0;
        let p = (a[0] as u128) * (a[1] as u128) + (carry as u128);
        cross[1] = p as u64;
        carry = (p >> 64) as u64;
        let p = (a[0] as u128) * (a[2] as u128) + (carry as u128);
        cross[2] = p as u64;
        carry = (p >> 64) as u64;
        let p = (a[0] as u128) * (a[3] as u128) + (carry as u128);
        cross[3] = p as u64;
        cross[4] = (p >> 64) as u64;

        carry = 0;
        let p = (a[1] as u128) * (a[2] as u128) + (cross[3] as u128) + (carry as u128);
        cross[3] = p as u64;
        carry = (p >> 64) as u64;
        let p = (a[1] as u128) * (a[3] as u128) + (cross[4] as u128) + (carry as u128);
        cross[4] = p as u64;
        cross[5] = (p >> 64) as u64;

        let p = (a[2] as u128) * (a[3] as u128) + (cross[5] as u128);
        cross[5] = p as u64;
        cross[6] = (p >> 64) as u64;

        // Step 2: Double the cross products (left shift by 1 bit).
        let mut c = 0u64;
        let mut k = 0;
        while k < 8 {
            let new_c = cross[k] >> 63;
            cross[k] = (cross[k] << 1) | c;
            c = new_c;
            k += 1;
        }

        // Step 3: Add diagonal products with carry propagation.
        // Each diagonal a[i]^2 contributes a 128-bit value at position 2*i.
        let mut r = [0u64; 8];
        let mut carry: u64 = 0;

        let mut i = 0;
        while i < 4 {
            let sq = (a[i] as u128) * (a[i] as u128);
            let k = 2 * i;

            let sum = (cross[k] as u128) + (sq as u64 as u128) + (carry as u128);
            r[k] = sum as u64;

            let sum = (cross[k + 1] as u128) + (sq >> 64) + (sum >> 64);
            r[k + 1] = sum as u64;
            carry = (sum >> 64) as u64;

            i += 1;
        }

        U512(r)
    }
}

#[cfg(test)]
mod ai_tests {
    use super::*;
    /// 7^2 = 49.
    #[test]
    fn small_square() {
        let a = U256::from_be_limbs([0, 0, 0, 7]);
        assert_eq!(a.square(), U512([49, 0, 0, 0, 0, 0, 0, 0]));
    }

    /// 0^2 = 0.
    #[test]
    fn square_zero() {
        assert_eq!(U256::ZERO.square(), U512([0; 8]));
    }

    /// 1^2 = 1.
    #[test]
    fn square_one() {
        assert_eq!(U256::ONE.square(), U512([1, 0, 0, 0, 0, 0, 0, 0]));
    }

    /// (2^64)^2 = 2^128.
    #[test]
    fn cross_limb_square() {
        let a = U256::from_be_limbs([0, 0, 1, 0]);
        assert_eq!(a.square(), U512([0, 0, 1, 0, 0, 0, 0, 0]));
    }

    /// (2^64 + 1)^2 = 2^128 + 2^65 + 1.
    #[test]
    fn cross_product_square() {
        let a = U256::from_be_limbs([0, 0, 1, 1]);
        // 2^128 + 2^65 + 1 = [0,0,0,0, 0, 1, 2, 1] in BE u64 limbs
        assert_eq!(a.square(), U512([1, 2, 1, 0, 0, 0, 0, 0]));
    }

    /// Max * Max = (2^256 - 1)^2 = 2^512 - 2^257 + 1.
    #[test]
    fn max_square() {
        let max = U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]);
        assert_eq!(
            max.square(),
            U512([
                1,
                0,
                0,
                0,
                0xFFFFFFFFFFFFFFFE,
                0xFFFFFFFFFFFFFFFF,
                0xFFFFFFFFFFFFFFFF,
                0xFFFFFFFFFFFFFFFF,
            ]),
        );
    }

    /// Squaring matches `a * a` for various inputs.
    #[test]
    fn matches_mul() {
        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([
                0xFFFFFFFFFFFFFFFF,
                0xFFFFFFFFFFFFFFFF,
                0xFFFFFFFFFFFFFFFF,
                0xFFFFFFFEFFFFFC2E,
            ]),
        ];
        for a in &cases {
            assert_eq!(
                a.square(),
                *a * *a,
                "mismatch for {a:?}",
            );
        }
    }
}
