use anyhow::Result;
use ethereum_types::U256;
use rand::{thread_rng, Rng};

use crate::cpu::kernel::aggregator::KERNEL;
use crate::cpu::kernel::interpreter::run_interpreter;

// TODO: 107 is hardcoded as a dummy prime for testing
// should be changed to the proper implementation prime
// once the run_{add, mul, sub}fp254 fns are implemented
const P254: u32 = 107;

fn add_fp(x: u32, y: u32) -> u32 {
    (x + y) % P254
}

fn add3_fp(x: u32, y: u32, z: u32) -> u32 {
    (x + y + z) % P254
}

fn mul_fp(x: u32, y: u32) -> u32 {
    (x * y) % P254
}

fn sub_fp(x: u32, y: u32) -> u32 {
    (P254 + x - y) % P254
}

fn add_fp2(a: [u32; 2], b: [u32; 2]) -> [u32; 2] {
    let [a, a_] = a;
    let [b, b_] = b;
    [add_fp(a, b), add_fp(a_, b_)]
}

fn add3_fp2(a: [u32; 2], b: [u32; 2], c: [u32; 2]) -> [u32; 2] {
    let [a, a_] = a;
    let [b, b_] = b;
    let [c, c_] = c;
    [add3_fp(a, b, c), add3_fp(a_, b_, c_)]
}

fn sub_fp2(a: [u32; 2], b: [u32; 2]) -> [u32; 2] {
    let [a, a_] = a;
    let [b, b_] = b;
    [sub_fp(a, b), sub_fp(a_, b_)]
}

fn mul_fp2(a: [u32; 2], b: [u32; 2]) -> [u32; 2] {
    let [a, a_] = a;
    let [b, b_] = b;
    [
        sub_fp(mul_fp(a, b), mul_fp(a_, b_)),
        add_fp(mul_fp(a, b_), mul_fp(a_, b)),
    ]
}

fn i9(a: [u32; 2]) -> [u32; 2] {
    let [a, a_] = a;
    [sub_fp(mul_fp(9, a), a_), add_fp(a, mul_fp(9, a_))]
}

fn add_fp6(c: [[u32; 2]; 3], d: [[u32; 2]; 3]) -> [[u32; 2]; 3] {
    let [c0, c1, c2] = c;
    let [d0, d1, d2] = d;

    let e0 = add_fp2(c0, d0);
    let e1 = add_fp2(c1, d1);
    let e2 = add_fp2(c2, d2);
    [e0, e1, e2]
}

fn sub_fp6(c: [[u32; 2]; 3], d: [[u32; 2]; 3]) -> [[u32; 2]; 3] {
    let [c0, c1, c2] = c;
    let [d0, d1, d2] = d;

    let e0 = sub_fp2(c0, d0);
    let e1 = sub_fp2(c1, d1);
    let e2 = sub_fp2(c2, d2);
    [e0, e1, e2]
}

fn mul_fp6(c: [[u32; 2]; 3], d: [[u32; 2]; 3]) -> [[u32; 2]; 3] {
    let [c0, c1, c2] = c;
    let [d0, d1, d2] = d;

    let c0d0 = mul_fp2(c0, d0);
    let c0d1 = mul_fp2(c0, d1);
    let c0d2 = mul_fp2(c0, d2);
    let c1d0 = mul_fp2(c1, d0);
    let c1d1 = mul_fp2(c1, d1);
    let c1d2 = mul_fp2(c1, d2);
    let c2d0 = mul_fp2(c2, d0);
    let c2d1 = mul_fp2(c2, d1);
    let c2d2 = mul_fp2(c2, d2);
    let cd12 = add_fp2(c1d2, c2d1);

    [
        add_fp2(c0d0, i9(cd12)),
        add3_fp2(c0d1, c1d0, i9(c2d2)),
        add3_fp2(c0d2, c1d1, c2d0),
    ]
}

fn sh(c: [[u32; 2]; 3]) -> [[u32; 2]; 3] {
    let [c0, c1, c2] = c;
    [i9(c2), c0, c1]
}

fn sparse_embed(x: [u32; 5]) -> [[[u32; 2]; 3]; 2] {
    let [g0, g1, g1_, g2, g2_] = x;
    [[[g0, 0], [g1, g1_], [0, 0]], [[0, 0], [g2, g2_], [0, 0]]]
}

fn mul_fp12(f: [[[u32; 2]; 3]; 2], g: [[[u32; 2]; 3]; 2]) -> [[[u32; 2]; 3]; 2] {
    let [f0, f1] = f;
    let [g0, g1] = g;

    let h0 = mul_fp6(f0, g0);
    let h1 = mul_fp6(f1, g1);
    let h01 = mul_fp6(add_fp6(f0, f1), add_fp6(g0, g1));
    [add_fp6(h0, sh(h1)), sub_fp6(h01, add_fp6(h0, h1))]
}

fn gen_fp6() -> [[u32; 2]; 3] {
    let mut rng = thread_rng();
    [
        [rng.gen_range(0..P254), rng.gen_range(0..P254)],
        [rng.gen_range(0..P254), rng.gen_range(0..P254)],
        [rng.gen_range(0..P254), rng.gen_range(0..P254)],
    ]
}

fn gen_fp12_sparse() -> [[[u32; 2]; 3]; 2] {
    let mut rng = thread_rng();
    sparse_embed([
        rng.gen_range(0..P254),
        rng.gen_range(0..P254),
        rng.gen_range(0..P254),
        rng.gen_range(0..P254),
        rng.gen_range(0..P254),
    ])
}

fn as_stack(xs: Vec<u32>) -> Vec<U256> {
    xs.iter().map(|&x| U256::from(x)).rev().collect()
}

#[test]
fn test_fp6() -> Result<()> {
    let c = gen_fp6();
    let d = gen_fp6();

    let mut input: Vec<u32> = [c, d].into_iter().flatten().flatten().collect();
    input.push(0xdeadbeef);

    let initial_offset = KERNEL.global_labels["mul_fp6"];
    let initial_stack: Vec<U256> = as_stack(input);
    let final_stack: Vec<U256> = run_interpreter(initial_offset, initial_stack)?
        .stack()
        .to_vec();

    let output: Vec<u32> = mul_fp6(c, d).into_iter().flatten().collect();
    let expected = as_stack(output);

    assert_eq!(final_stack, expected);

    Ok(())
}

fn make_initial_stack(
    in1: u32,
    in2: u32,
    out: u32,
    f0: [[u32; 2]; 3],
    f1: [[u32; 2]; 3],
    g0: [[u32; 2]; 3],
    g1: [[u32; 2]; 3],
) -> Vec<U256> {
    // stack: in0, f, in0', f', in1, g, in1', g', in1, out, in0, out
    let f0: Vec<u32> = f0.into_iter().flatten().collect();
    let f1: Vec<u32> = f1.into_iter().flatten().collect();
    let g0: Vec<u32> = g0.into_iter().flatten().collect();
    let g1: Vec<u32> = g1.into_iter().flatten().collect();

    let mut input = f0;
    input.extend(vec![in1]);
    input.extend(f1);
    input.extend(g0);
    input.extend(vec![in2]);
    input.extend(g1);
    input.extend(vec![in2, out, in1, out]);

    as_stack(input)
}

// #[test]
fn test_fp12() -> Result<()> {
    let in1 = 64;
    let in2 = 76;
    let out = 88;

    let f0 = gen_fp6();
    let f1 = gen_fp6();
    let g0 = gen_fp6();
    let g1 = gen_fp6();

    let initial_offset = KERNEL.global_labels["test_mul_fp12"];
    let initial_stack: Vec<U256> = make_initial_stack(in1, in2, out, f0, f1, g0, g1);
    let final_stack: Vec<U256> = run_interpreter(initial_offset, initial_stack)?
        .stack()
        .to_vec();

    let mut output: Vec<u32> = mul_fp12([f0, f1], [g0, g1])
        .into_iter()
        .flatten()
        .flatten()
        .collect();
    output.extend(vec![out]);
    let expected = as_stack(output);

    assert_eq!(final_stack, expected);

    Ok(())
}

// #[test]
fn test_fp12_sparse() -> Result<()> {
    let in1 = 64;
    let in2 = 76;
    let out = 88;

    let f0 = gen_fp6();
    let f1 = gen_fp6();
    let [g0, g1] = gen_fp12_sparse();

    let initial_offset = KERNEL.global_labels["test_mul_fp12"];
    let initial_stack: Vec<U256> = make_initial_stack(in1, in2, out, f0, f1, g0, g1);
    let final_stack: Vec<U256> = run_interpreter(initial_offset, initial_stack)?
        .stack()
        .to_vec();

    let mut output: Vec<u32> = mul_fp12([f0, f1], [g0, g1])
        .into_iter()
        .flatten()
        .flatten()
        .collect();
    output.extend(vec![out]);
    let expected = as_stack(output);

    assert_eq!(final_stack, expected);

    Ok(())
}

#[test]
fn test_fp12_square() -> Result<()> {
    let in1 = 64;
    let in2 = 76;
    let out = 88;

    let f0 = gen_fp6();
    let f1 = gen_fp6();

    let initial_offset = KERNEL.global_labels["test_mul_fp12"];
    let initial_stack: Vec<U256> = make_initial_stack(in1, in2, out, f0, f1, f0, f1);
    let final_stack: Vec<U256> = run_interpreter(initial_offset, initial_stack)?
        .stack()
        .to_vec();

    let mut output: Vec<u32> = mul_fp12([f0, f1], [f0, f1])
        .into_iter()
        .flatten()
        .flatten()
        .collect();
    output.extend(vec![out]);
    let expected = as_stack(output);

    assert_eq!(final_stack, expected);

    Ok(())
}