//! Montgomery batch inversion for computing multiple modular inverses
//! efficiently using a single expensive `mod_inv` call.

use super::U256Batch;
use crate::u256::U256;

impl<const N: usize> U256Batch<N> {
    /// Computes the modular inverse of each element using Montgomery's batch
    /// inversion trick: one `mod_inv` call plus `3(N-1)` modular
    /// multiplications, amortizing the cost of inversion across the batch.
    ///
    /// The algorithm works in three phases:
    /// 1. **Forward sweep**: build prefix products `out[i] = v[0] * v[1] * ... * v[i]`.
    /// 2. **Single inversion**: invert the total product `out[N-1]`.
    /// 3. **Backward sweep**: peel off individual inverses by multiplying
    ///    the running inverse by the corresponding prefix product.
    ///
    /// Returns `None` if any element is zero or not invertible modulo
    /// `modulus`. Returns the raw `[U256; N]` array rather than a
    /// `U256Batch` so the caller decides how to use the results.
    ///
    /// For the vanity address mining hot loop, this converts 64 Jacobian
    /// Z-coordinates to affine using ~48ns/inversion (amortized) instead of
    /// ~785ns each.
    ///
    /// # Examples
    ///
    /// ```
    /// use cnfy_uint::u256_batch::U256Batch;
    /// use cnfy_uint::u256::U256;
    ///
    /// let p = U256::from_be_limbs([
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    /// ]);
    /// let batch = U256Batch([
    ///     U256::from_be_limbs([0, 0, 0, 3]),
    ///     U256::from_be_limbs([0, 0, 0, 7]),
    /// ]);
    /// let inverses = batch.mod_inv(&p).unwrap();
    ///
    /// // Verify each result: v[i] * inv[i] = 1 (mod p)
    /// for i in 0..2 {
    ///     assert_eq!(batch.0[i].mul_mod(&inverses[i], &p), U256::ONE);
    /// }
    /// ```
    pub fn mod_inv(&self, modulus: &U256) -> Option<[U256; N]> {
        if N == 0 {
            return Some([U256::ZERO; N]);
        }

        let values = &self.0;
        let mut out = [U256::ZERO; N];

        // Phase 1: forward sweep — prefix products stored in `out`.
        out[0] = values[0];
        for i in 1..N {
            out[i] = out[i - 1].mul_mod(&values[i], modulus);
        }

        // Phase 2: single inversion of the total product.
        let mut inv = out[N - 1].mod_inv(modulus)?;

        // Phase 3: backward sweep — extract individual inverses.
        for i in (1..N).rev() {
            // out[i] = inv * prefix[i-1] = inverse of values[i]
            out[i] = inv.mul_mod(&out[i - 1], modulus);
            // Update running inverse: multiply by values[i] to peel it off
            inv = inv.mul_mod(&values[i], modulus);
        }
        out[0] = inv;

        Some(out)
    }
}

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

    /// Empty batch returns an empty array.
    #[test]
    fn empty_batch() {
        let batch = U256Batch::<0>::zeroed();
        let result = batch.mod_inv(&P);
        assert!(result.is_some());
        assert_eq!(result.unwrap().len(), 0);
    }

    /// Single element matches individual mod_inv.
    #[test]
    fn single_element() {
        let batch = U256Batch([U256::from_be_limbs([0, 0, 0, 7])]);
        let result = batch.mod_inv(&P).unwrap();

        let expected = U256::from_be_limbs([0, 0, 0, 7]).mod_inv(&P).unwrap();
        assert_eq!(result[0], expected);
    }

    /// Three small values: verify each inverse individually.
    #[test]
    fn three_small_values() {
        let batch = U256Batch([
            U256::from_be_limbs([0, 0, 0, 2]),
            U256::from_be_limbs([0, 0, 0, 3]),
            U256::from_be_limbs([0, 0, 0, 5]),
        ]);
        let result = batch.mod_inv(&P).unwrap();

        for i in 0..3 {
            assert_eq!(
                batch.0[i].mul_mod(&result[i], &P),
                U256::ONE,
                "inverse failed for index {i}"
            );
        }
    }

    /// Batch of 8 values matches individual mod_inv results.
    #[test]
    fn batch_vs_individual_eight() {
        let batch = U256Batch([
            U256::from_be_limbs([0, 0, 0, 2]),
            U256::from_be_limbs([0, 0, 0, 3]),
            U256::from_be_limbs([0, 0, 0, 5]),
            U256::from_be_limbs([0, 0, 0, 7]),
            U256::from_be_limbs([0, 0, 0, 11]),
            U256::from_be_limbs([0, 0, 0, 13]),
            U256::from_be_limbs([0, 0, 0, 17]),
            U256::from_be_limbs([0, 0, 0, 19]),
        ]);
        let result = batch.mod_inv(&P).unwrap();

        for i in 0..8 {
            let individual = batch.0[i].mod_inv(&P).unwrap();
            assert_eq!(
                result[i], individual,
                "mismatch at index {i}: batch={:?}, individual={:?}",
                result[i], individual
            );
        }
    }

    /// Full-width U256 values produce correct inverses.
    #[test]
    fn full_width_values() {
        const GX: U256 = U256::from_be_limbs([0x79BE667EF9DCBBAC, 0x55A06295CE870B07, 0x029BFCDB2DCE28D9, 0x59F2815B16F81798]);

        let gy = U256::from_be_limbs([
            0x483ADA7726A3C465,
            0x5DA4FBFC0E1108A8,
            0xFD17B448A6855419,
            0x9C47D08FFB10D4B8,
        ]);

        let batch = U256Batch([GX, gy]);
        let result = batch.mod_inv(&P).unwrap();

        assert_eq!(GX.mul_mod(&result[0], &P), U256::ONE);
        assert_eq!(gy.mul_mod(&result[1], &P), U256::ONE);
    }

    /// Returns None when an element is zero.
    #[test]
    fn none_on_zero_element() {
        let batch = U256Batch([U256::from_be_limbs([0, 0, 0, 3]), U256::ZERO]);
        assert!(batch.mod_inv(&P).is_none());
    }

    /// Returns None when an element equals the modulus (not invertible).
    #[test]
    fn none_on_modulus_element() {
        let batch = U256Batch([U256::from_be_limbs([0, 0, 0, 3]), P]);
        assert!(batch.mod_inv(&P).is_none());
    }
}
