//! Full-width `(f, g)` matrix application for [`DivstepMatrix`].
use super::DivstepMatrix;
use crate::u256::U256;
use crate::u320::U320;

impl DivstepMatrix {
    /// Applies the accumulated transition matrix to full-width signed 256-bit
    /// values `f` and `g`, then right-shifts by 62 bits.
    ///
    /// Computes:
    /// ```text
    /// f' = (u*f + v*g) >> 62
    /// g' = (q*f + r*g) >> 62
    /// ```
    ///
    /// Products are computed at 320-bit width via `U256 * u64 -> U320` and
    /// combined with signed addition. The divstep algorithm guarantees both
    /// results are divisible by `2^62` and fit in 256 bits after the shift.
    ///
    /// Signed representation: `(magnitude, negative)` where `magnitude` is the
    /// absolute value and `negative` is true when the value is negative.
    pub(crate) fn apply_fg(
        &self,
        f: &U256,
        f_neg: bool,
        g: &U256,
        g_neg: bool,
    ) -> (U256, bool, U256, bool) {
        let (new_f, new_f_neg) = Self::signed_linear_combo_320(self.u, self.v, f, f_neg, g, g_neg);
        let (new_g, new_g_neg) = Self::signed_linear_combo_320(self.q, self.r, f, f_neg, g, g_neg);

        // Right-shift by 62 bits and extract to U256
        let new_f = Self::shr62_320(&new_f);
        let new_g = Self::shr62_320(&new_g);

        (new_f, new_f_neg, new_g, new_g_neg)
    }

    /// Computes `|a * x + b * y|` as a U320 with sign, where `a`, `b` are
    /// signed i64 and `x`, `y` are signed 256-bit values.
    fn signed_linear_combo_320(
        a: i64,
        b: i64,
        x: &U256,
        x_neg: bool,
        y: &U256,
        y_neg: bool,
    ) -> (U320, bool) {
        // Compute |a| * |x| as U320
        let a_neg = a < 0;
        let a_abs = a.unsigned_abs();
        let ax = *x * a_abs;
        let ax_neg = a_neg ^ x_neg;

        // Compute |b| * |y| as U320
        let b_neg = b < 0;
        let b_abs = b.unsigned_abs();
        let by = *y * b_abs;
        let by_neg = b_neg ^ y_neg;

        // Add with signs at U320 width
        Self::signed_add_320(&ax, ax_neg, &by, by_neg)
    }

    /// Adds two signed U320 values, returning `(|result|, result_negative)`.
    fn signed_add_320(a: &U320, a_neg: bool, b: &U320, b_neg: bool) -> (U320, bool) {
        if a_neg == b_neg {
            let (sum, _) = a.overflowing_add(b);
            (sum, a_neg)
        } else if a >= b {
            let (diff, _) = a.overflowing_sub(b);
            (diff, a_neg)
        } else {
            let (diff, _) = b.overflowing_sub(a);
            (diff, b_neg)
        }
    }

    /// Right-shifts a U320 by 62 bits and returns the low 256 bits as U256.
    ///
    /// The divstep algorithm guarantees the shifted result fits in 256 bits.
    fn shr62_320(v: &U320) -> U256 {
        let [w0, w1, w2, w3, w4] = v.0;
        U256([
            (w0 >> 62) | (w1 << 2),
            (w1 >> 62) | (w2 << 2),
            (w2 >> 62) | (w3 << 2),
            (w3 >> 62) | (w4 << 2),
        ])
    }
}

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

    /// shr62 of 2^62 yields 1.
    #[test]
    fn shr62_basic() {
        let v = U320::from_be_limbs([0, 0, 0, 0, 1u64 << 62]);
        let shifted = DivstepMatrix::shr62_320(&v);
        assert_eq!(shifted, U256::ONE);
    }

    /// shr62 across U320 limb boundaries.
    #[test]
    fn shr62_cross_limb() {
        // LE: [0, 1, 0, 0, 0] = value has 1 in limb index 1
        let v = U320::from_be_limbs([0, 0, 0, 1, 0]);
        let shifted = DivstepMatrix::shr62_320(&v);
        // After >>62: limb[0] = (0>>62)|(1<<2) = 4, rest 0
        assert_eq!(shifted, U256([4, 0, 0, 0]));
    }

    /// shr62 with a value in the 5th limb.
    #[test]
    fn shr62_fifth_limb() {
        // LE: [0, 0, 0, 0, 1] = 2^256
        let v = U320::from_be_limbs([1, 0, 0, 0, 0]);
        let shifted = DivstepMatrix::shr62_320(&v);
        // After >>62: limb[3] = (0>>62)|(1<<2) = 4
        assert_eq!(shifted, U256([0, 0, 0, 4]));
    }

    /// signed_add_320 with same signs adds magnitudes.
    #[test]
    fn signed_add_same_sign() {
        let a = U320::from_be_limbs([0, 0, 0, 0, 5]);
        let b = U320::from_be_limbs([0, 0, 0, 0, 3]);
        let (result, neg) = DivstepMatrix::signed_add_320(&a, false, &b, false);
        assert_eq!(result, U320::from_be_limbs([0, 0, 0, 0, 8]));
        assert!(!neg);
    }

    /// signed_add_320 with different signs subtracts.
    #[test]
    fn signed_add_different_sign() {
        let a = U320::from_be_limbs([0, 0, 0, 0, 5]);
        let b = U320::from_be_limbs([0, 0, 0, 0, 3]);
        let (result, neg) = DivstepMatrix::signed_add_320(&a, false, &b, true);
        assert_eq!(result, U320::from_be_limbs([0, 0, 0, 0, 2]));
        assert!(!neg);
    }
}
