+

use super::primefield;
use std::u128;

pub fn multiply_mod_prime(a: primefield::Element, b: primefield::Element) -> primefield::Element {
    const MASK_48: u128 = 0xFFFFFFFFFFFF;
    const MASK_52: u128 = 0xFFFFFFFFFFFFF;
    const MASK_64: u128 = 0xFFFFFFFFFFFFFFFF;
    const R: u128 = 0x1000003D10;

    let a0 = a.limbs[0] as u128;
    let a1 = a.limbs[1] as u128;
    let a2 = a.limbs[2] as u128;
    let a3 = a.limbs[3] as u128;
    let a4 = a.limbs[4] as u128;

    let b0 = b.limbs[0] as u128;
    let b1 = b.limbs[1] as u128;
    let b2 = b.limbs[2] as u128;
    let b3 = b.limbs[3] as u128;
    let b4 = b.limbs[4] as u128;

    //  [... a b c] is a shorthand for ... + a<<104 + b<<52 + c<<0 mod n.
    //  for 0 <= x <= 4, px is a shorthand for sum(a[i]*b[x-i], i=0..x).
    //  for 4 <= x <= 8, px is a shorthand for sum(a[i]*b[x-i], i=(x-4)..4)
    //  Note that [x 0 0 0 0 0] = [x*R].

    let mut d = a0 * b3;
    d += a1 * b2;
    d += a2 * b1;
    d += a3 * b0;
    assert!(d >> 114 == 0);
    // [d 0 0 0] = [p3 0 0 0]

    let mut c = a4 * b4;
    assert!(c >> 112 == 0);
    // [c 0 0 0 0 d 0 0 0] = [p8 0 0 0 0 p3 0 0 0]

    d += (c & MASK_64) * R;
    c >>= 64;
    assert!(d >> 115 == 0);
    assert!(c >> 48 == 0);
    // [(c<<12) 0 0 0 0 0 d 0 0 0] = [p8 0 0 0 0 p3 0 0 0]

    let t3 = d & MASK_52;
    d >>= 52;
    assert!(t3 >> 52 == 0);
    assert!(d >> 63 == 0);
    // [(c<<12) 0 0 0 0 d t3 0 0 0] = [p8 0 0 0 0 p3 0 0 0]

    d += a0 * b4;
    d += a1 * b3;
    d += a2 * b2;
    d += a3 * b1;
    d += a4 * b0;
    assert!(d >> 115 == 0);
    // [(c<<12) 0 0 0 0 d t3 0 0 0] = [p8 0 0 0 p4 p3 0 0 0]

    d += (c & MASK_64) * (R << 12);
    assert!(d >> 116 == 0);
    // [d t3 0 0 0] = [p8 0 0 0 p4 p3 0 0 0]

    let mut t4 = d & MASK_52;
    d >>= 52;
    assert!(t4 >> 52 == 0);
    assert!(d >> 64 == 0);
    // [d t4 t3 0 0 0] = [p8 0 0 0 p4 p3 0 0 0]

    let tx = t4 >> 48;
    t4 &= MASK_48;
    assert!(tx >> 4 == 0);
    assert!(t4 >> 48 == 0);
    // [d (t4 + (tx << 48)) t3 0 0 0] = [p8 0 0 0 p4 p3 0 0 0] */
    //
    c = a0 * b0;
    assert!(c >> 112 == 0);
    // [d t4+(tx<<48) t3 0 0 c] = [p8 0 0 0 p4 p3 0 0 p0]

    d += a1 * b4;
    d += a2 * b3;
    d += a3 * b2;
    d += a4 * b1;
    assert!(d >> 115 == 0);
    // [d (t4 + (tx << 48)) t3 0 0 c] = [p8 0 0 p5 p4 p3 0 0 p0]

    let mut u0 = d & MASK_52;
    d >>= 52;
    assert!(u0 >> 52 == 0);
    assert!(d >> 63 == 0);
    // [d u0 (t4 + (tx << 48)) t3 0 0 c] = [p8 0 0 p5 p4 p3 0 0 p0]
    // [d 0 (t4 + (tx << 48) + (u0 << 52)) t3 0 0 c] = [p8 0 0 p5 p4 p3 0 0 p0]

    u0 = (u0 << 4) | tx;
    assert!(u0 >> 56 == 0);
    // [d 0 t4+(u0<<48) t3 0 0 c] = [p8 0 0 p5 p4 p3 0 0 p0]

    c += u0 * (R >> 4);
    assert!(c >> 115 == 0);
    // [d 0 t4 t3 0 0 c] = [p8 0 0 p5 p4 p3 0 0 p0]

    let r0 = c & MASK_52;
    c >>= 52;
    assert!(r0 >> 52 == 0);
    assert!(c >> 61 == 0);
    // [d 0 t4 t3 0 c r0] = [p8 0 0 p5 p4 p3 0 0 p0]

    c += a0 * b1;
    c += a1 * b0;
    assert!(c >> 114 == 0);
    // [d 0 t4 t3 0 c r0] = [p8 0 0 p5 p4 p3 0 p1 p0]

    d += a2 * b4;
    d += a3 * b3;
    d += a4 * b2;
    assert!(d >> 114 == 0);
    // [d 0 t4 t3 0 c r0] = [p8 0 p6 p5 p4 p3 0 p1 p0]

    c += (d & MASK_52) * R;
    d >>= 52;
    assert!(c >> 115 == 0);
    assert!(d >> 62 == 0);
    // [d 0 0 t4 t3 0 c r0] = [p8 0 p6 p5 p4 p3 0 p1 p0]

    let r1 = c & MASK_52;
    c >>= 52;
    assert!(r1 >> 52 == 0);
    assert!(c >> 63 == 0);
    // [d 0 0 t4 t3 c r1 r0] = [p8 0 p6 p5 p4 p3 0 p1 p0]

    c += a0 * b2;
    c += a1 * b1;
    c += a2 * b0;
    assert!(c >> 114 == 0);
    // [d 0 0 t4 t3 c r1 r0] = [p8 0 p6 p5 p4 p3 p2 p1 p0]

    d += a3 * b4;
    d += a4 * b3;
    assert!(d >> 114 == 0);
    // [d 0 0 t4 t3 c t1 r0] = [p8 p7 p6 p5 p4 p3 p2 p1 p0]

    c += R * (d & MASK_64);
    d >>= 64;
    assert!(c >> 115 == 0);
    assert!(d >> 50 == 0);
    // [(d<<12) 0 0 0 t4 t3 c r1 r0] = [p8 p7 p6 p5 p4 p3 p2 p1 p0]

    let r2 = c & MASK_52;
    c >>= 52;
    assert!(r2 >> 52 == 0);
    assert!(c >> 63 == 0);
    // [(d<<12) 0 0 0 t4 (t3 + c) r2 r1 r0] = [p8 p7 p6 p5 p4 p3 p2 p1 p0]

    c += (d & MASK_64) * (R << 12);
    c += t3;
    assert!(c >> 100 == 0);
    // [t4 c r2 r1 r0] = [p8 p7 p6 p5 p4 p3 p2 p1 p0]

    let r3 = c & MASK_52;
    c >>= 52;
    assert!(r3 >> 52 == 0);
    assert!(c >> 48 == 0);
    // [t4+c r3 r2 r1 r0] = [p8 p7 p6 p5 p4 p3 p2 p1 p0]

    let r4 = (c & MASK_64) + t4;
    assert!(r4 >> 49 == 0);
    // [r4 r3 r2 r1 r0] = [p8 p7 p6 p5 p4 p3 p2 p1 p0]

    let limbs = [r0 as u64, r1 as u64, r2 as u64, r3 as u64, r4 as u64];
    let magnitude = 0;

    primefield::Element { limbs, magnitude }
}

pub fn invert_mod_prime(x: primefield::Element) -> primefield::Element {
    let mut delta = 1;
    let mut f = [
    0x3FFFFFFEFFFFFC2F,
    0x3FFFFFFFFFFFFFFF,
    0x3FFFFFFFFFFFFFFF,
    0x3FFFFFFFFFFFFFFF,
    0xFF,
];
    let mut g = field_element_to_signed_62(x);
    let mut d = primefield::Element::ZERO;
    let mut e = primefield::Element::ONE;

    for _ in 0..12 {
        let (new_delta, u, v, q, r) = matrix_62(delta, f[0] as i64, g[0] as i64);
        delta = new_delta;

        update_fg_62(&mut f, &mut g, u, v, q, r);

        d = d * i64_to_field_element(u) + e * i64_to_field_element(v);
        e = d * i64_to_field_element(q) + e * i64_to_field_element(r);
    }

    signed_62_to_field_element(f) * d * primefield::Element::INVERSE_2_POW_744
}

fn matrix_62(mut delta: i64, mut f: i64, mut g: i64) -> (i64, i64, i64, i64, i64) {
    //Compute delta and transition matrix t after N divsteps (multiplied by 2^N).

    let (mut u, mut v, mut q, mut r) = (1, 0, 0, 1); // start with identity matrix

    for _ in 0..62 {
        assert!(f & 1 == 1);
        if delta > 0 && g & 1 == 1 {
            (delta, f, g, u, v, q, r) = (1 - delta, g, (g - f) >> 1, 2 * q, 2 * r, q - u, r - v);
        } else if g & 1 == 1 {
            (delta, f, g, u, v, q, r) = (1 + delta, f, (g + f) >> 1, 2 * u, 2 * v, q + u, r + v);
        } else {
            (delta, f, g, u, v, q, r) = (1 + delta, f, g >> 1, 2 * u, 2 * v, q, r);
        }
    }

    (delta, u, v, q, r)
}

fn update_fg_62(f: &mut [i128; 5], g: &mut [i128; 5], u: i64, v: i64, q: i64, r: i64) {
    const MASK_62: i128 = i128::MAX >> 65;

    let u = u as i128;
    let v = v as i128;
    let q = q as i128;
    let r = r as i128;

    // Start computing t*[f,g].
    let mut cf = u * f[0] + v * g[0];
    let mut cg = q * f[0] + r * g[0];

    // Verify that the bottom 62 bits of the result are zero, and then throw them away.
    assert!(cf & MASK_62 == 0);
    assert!(cg & MASK_62 == 0);
    cf >>= 62;
    cg >>= 62;

    for i in 0..4 {
        // Compute limb i + 1 of t*[f,g], and store it as output limb i shifting down by 62 bits
        cf += u * f[i + 1] + v * g[i + 1];
        cg += q * f[i + 1] + r * g[i + 1];

        f[i] = cf & MASK_62;
        g[i] = cg & MASK_62;

        cf >>= 62;
        cg >>= 62;
    }

    // What remains is limb 5 of t*[f,g]; store it as output limb 4.
    f[4] = cf & MASK_62;
    g[4] = cg & MASK_62;
}

fn field_element_to_signed_62(x: primefield::Element) -> [i128; 5] {
    const MASK_62: u64 = u64::MAX >> 2;

    let x = x.normalize();

    let x0 = x.limbs[0];
    let x1 = x.limbs[1];
    let x2 = x.limbs[2];
    let x3 = x.limbs[3];
    let x4 = x.limbs[4];

    let r0 = (x0 | (x1 << 52)) & MASK_62;
    let r1 = (x1 >> 10 | x2 << 42) & MASK_62;
    let r2 = (x2 >> 20 | x3 << 32) & MASK_62;
    let r3 = (x3 >> 30 | x4 << 22) & MASK_62;
    let r4 = (x4 >> 40) & MASK_62;

    [r0 as i128, r1 as i128, r2 as i128, r3 as i128, r4 as i128]
}

fn signed_62_to_field_element(x: [i128; 5]) -> primefield::Element {
    const MASK_48: u64 = u64::MAX >> 16;
    const MASK_52: u64 = u64::MAX >> 12;

    let n0 = x[0].is_negative();
    let n1 = x[1].is_negative();
    let n2 = x[2].is_negative();
    let n3 = x[3].is_negative();
    let n4 = x[4].is_negative();

    let x0 = x[0] as u64;
    let x1 = x[1] as u64;
    let x2 = x[2] as u64;
    let x3 = x[3] as u64;
    let x4 = x[4] as u64;

    let e0 = primefield::Element::new([x0 & MASK_52, x0 >> 52, 0, 0, 0]);
    let e1 = primefield::Element::new([0, (x1 << 10) & MASK_52, x1 >> 42, 0, 0]);
    let e2 = primefield::Element::new([0, 0, (x2 << 20) & MASK_52, x2 >> 32, 0]);
    let e3 = primefield::Element::new([0, 0, 0, (x3 << 30) & MASK_52, x3 >> 22]);
    let e4 = primefield::Element::new([0, 0, 0, 0, (x4 << 40) & MASK_48]);

    let e0 = if n0 { -e0 } else { e0 };
    let e1 = if n1 { -e1 } else { e1 };
    let e2 = if n2 { -e2 } else { e2 };
    let e3 = if n3 { -e3 } else { e3 };
    let e4 = if n4 { -e4 } else { e4 };

    e0 + e1 + e2 + e3 + e4
}

fn i64_to_field_element(x: i64) -> primefield::Element {
    const MASK_52: u64 = u64::MAX >> 12;
    let n = x.is_negative();
    let x = x.abs() as u64;
    let e = primefield::Element::new([x & MASK_52, x >> 52, 0, 0, 0]);

    if n {
        -e
    } else {
        e
    }
}

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

    #[test]
    fn conversion() {
        println!("{:b}", i8::MIN.abs());
        let g = primefield::Element::GENERATOR_X;
        let signed = field_element_to_signed_62(g);
        let element = signed_62_to_field_element(signed);

        assert!(g == element);
    }
}