//! Validation of the dimension-4 theta arithmetic against the Phase 0 oracle.
//!
//! Two ground-truth files are used:
//!   * `../../sqisignhd-harness/test_vectors_l1.json` - the Phase 0 vectors
//!     (isogeny codomain theta null points), used for projective equality,
//!     normalisation and the Hadamard / middle-codomain relationships.
//!   * `tests/theta_arith_vectors.json` - the Phase 1 supplementary oracle
//!     (generic-point doubling and `act_point`), generated by
//!     `sqisignhd-harness/extract_theta_arith.py`.
//!
//! Every implemented operation is checked against the reference for all 5
//! Phase 0 vectors.

mod hd_common;
use hd_common::{fp2_eq, lambda, load, parse_coords, parse_node, F, PHASE0_VECTORS};

use sqisign_verify::hd::{act_point, hadamard, ThetaPointDim4, ThetaStructureDim4};

const THETA_ARITH_VECTORS: &str = concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/tests/theta_arith_vectors.json"
);

// 1. Projective equality (build and test this first)

#[test]
fn projective_eq_handbuilt() {
    // P = (1, 2, ..., 16)
    let p = ThetaPointDim4::new(core::array::from_fn(|k| F::from_small((k + 1) as u64)));
    // Scaling by any non-zero lambda gives an equal projective point.
    let scaled = p.scale(&lambda());
    assert!(p.projective_eq(&scaled), "P must equal lambda*P");
    assert!(scaled.projective_eq(&p), "equality is symmetric here");

    // Perturbing a single coordinate gives a different point.
    let mut coords: [F; 16] = core::array::from_fn(|k| F::from_small((k + 1) as u64));
    coords[7] = coords[7].add(&F::one());
    let perturbed = ThetaPointDim4::new(coords);
    assert!(!p.projective_eq(&perturbed), "perturbed point must differ");

    // A point with a leading zero coordinate: pivot must skip it.
    let mut z: [F; 16] = core::array::from_fn(|k| F::from_small((k + 1) as u64));
    z[0] = F::zero();
    let zp = ThetaPointDim4::new(z);
    assert_eq!(zp.pivot(), 1, "pivot skips leading zeros");
    assert!(zp.projective_eq(&zp.scale(&lambda())));
}

#[test]
fn projective_eq_against_vectors() {
    let doc = load(PHASE0_VECTORS);
    let vectors = doc["test_vectors"].as_array().unwrap();
    assert_eq!(vectors.len(), 5);
    for v in vectors {
        let s4 = &v["stage4_compute_hd"];
        let chain = s4["F1"]["chain_codomains"].as_array().unwrap();

        // A null point equals any non-zero rescaling of itself.
        let p = parse_coords(&chain[0]["coords"]);
        assert!(p.projective_eq(&p.scale(&lambda())));

        // Distinct codomains along the chain are distinct projective points.
        let q = parse_coords(&chain[1]["coords"]);
        assert!(!p.projective_eq(&q), "consecutive codomains must differ");
    }
}

// 2. Hadamard transform + middle-codomain relationships

/// Reference Hadamard: the O(16^2) definition, to cross-check the fast butterfly.
fn hadamard_slow(p: &[F; 16]) -> [F; 16] {
    core::array::from_fn(|chi| {
        let mut acc = F::zero();
        for (j, pj) in p.iter().enumerate() {
            if ((chi & j).count_ones() & 1) == 1 {
                acc = acc.sub(pj);
            } else {
                acc = acc.add(pj);
            }
        }
        acc
    })
}

#[test]
fn hadamard_fast_matches_slow() {
    let p: [F; 16] = core::array::from_fn(|k| {
        F::from_small((3 * k as u64 + 1) ^ 0xa5)
            .add(&F::i_element().mul(&F::from_small(k as u64 + 7)))
    });
    let fast = hadamard(&p);
    let slow = hadamard_slow(&p);
    for k in 0..16 {
        assert!(fp2_eq(&fast[k], &slow[k]), "hadamard mismatch at {k}");
    }
    // Hadamard is an involution up to the factor 16.
    let twice = hadamard(&fast);
    let sixteen = F::from_small(16);
    for k in 0..16 {
        assert!(fp2_eq(&twice[k], &p[k].mul(&sixteen)), "H^2 != 16*I at {k}");
    }
}

#[test]
fn hadamard_middle_codomain() {
    let doc = load(PHASE0_VECTORS);
    for v in doc["test_vectors"].as_array().unwrap() {
        let s4 = &v["stage4_compute_hd"];
        let s5 = &v["stage5_codomain_check"];

        let c1_zero = parse_node(&s5["C1_zero"]);
        let hc2_zero = parse_node(&s5["HC2_zero"]);

        // The verifier's own check: C1.zero() == C2.hadamard().zero().
        assert!(c1_zero.projective_eq(&hc2_zero), "C1.zero() == HC2.zero()");
        assert_eq!(s5["match"].as_bool(), Some(true));

        // C1.zero() equals the last F1 chain codomain (same structure).
        let f1 = s4["F1"]["chain_codomains"].as_array().unwrap();
        let last_f1 = parse_node(f1.last().unwrap());
        assert!(
            c1_zero.projective_eq(&last_f1),
            "C1.zero() == last F1 codomain"
        );

        // HC2.zero() == Hadamard(last F2_dual codomain), projectively.
        let f2 = s4["F2_dual"]["chain_codomains"].as_array().unwrap();
        let last_f2 = parse_node(f2.last().unwrap());
        let had = ThetaPointDim4::new(hadamard(last_f2.coords()));
        assert!(
            hc2_zero.projective_eq(&had),
            "HC2.zero() == Hadamard(last F2_dual codomain)"
        );
    }
}

// 3. Pivot normalisation matches the oracle's `normalized`/`pivot` fields

#[test]
fn normalize_matches_oracle() {
    let doc = load(PHASE0_VECTORS);
    for v in doc["test_vectors"].as_array().unwrap() {
        let s5 = &v["stage5_codomain_check"];
        for key in ["C1_zero", "HC2_zero"] {
            let node = &s5[key];
            let p = parse_node(node);
            let (pivot, normalized) = p.normalize();
            assert_eq!(
                pivot as u64,
                node["pivot"].as_u64().unwrap(),
                "{key} pivot mismatch"
            );
            let expected = parse_coords(&node["normalized"]);
            for k in 0..16 {
                assert!(
                    fp2_eq(&normalized.coords()[k], &expected.coords()[k]),
                    "{key} normalized coord {k} mismatch"
                );
            }
            // The pivot coordinate must be exactly 1.
            assert!(bool::from(normalized.coords()[pivot].ct_is_one()));
        }
    }
}

// 4. Doubling: double(null) == null, on every suitable codomain null point

#[test]
fn double_null_point_is_identity() {
    let doc = load(PHASE0_VECTORS);
    let mut suitable = 0usize;
    let mut total = 0usize;
    for v in doc["test_vectors"].as_array().unwrap() {
        let s4 = &v["stage4_compute_hd"];
        for chain_key in ["F1", "F2_dual"] {
            for node in s4[chain_key]["chain_codomains"].as_array().unwrap() {
                total += 1;
                let p = parse_node(node);
                let mut st = ThetaStructureDim4::new(p.clone());
                if !st.precompute() {
                    continue; // not suitable for the no-base-change doubling path
                }
                suitable += 1;
                let dbl = st.double(&p);
                assert!(
                    dbl.projective_eq(&p),
                    "double(null) must be projectively equal to null"
                );
            }
        }
    }
    // Sanity: the bulk of chain codomains support doubling.
    assert!(suitable > 0, "no suitable structures found");
    println!("double(null)==null verified on {suitable}/{total} codomain structures");
}

// 5. Generic-point doubling against the supplementary oracle

#[test]
fn doubling_oracle() {
    let doc = load(THETA_ARITH_VECTORS);
    assert_eq!(doc["suitable_doubling"].as_bool(), Some(true));

    let null = parse_coords(&doc["null_point"]);
    let mut st = ThetaStructureDim4::new(null.clone());
    assert!(st.precompute(), "oracle structure must support doubling");
    assert!(st.has_suitable_doubling());

    let cases = doc["cases"].as_array().unwrap();
    assert!(!cases.is_empty());
    for (n, case) in cases.iter().enumerate() {
        let p = parse_coords(&case["P"]);
        let expect_2p = parse_coords(&case["double"]);
        let expect_4p = parse_coords(&case["quad"]);

        let two_p = st.double(&p);
        assert!(two_p.projective_eq(&expect_2p), "case {n}: double(P) != 2P");

        let four_p = st.double_iter(&p, 2);
        assert!(
            four_p.projective_eq(&expect_4p),
            "case {n}: double_iter(P,2) != 4P"
        );

        // double_iter(P, 1) == double(P)
        assert!(st.double_iter(&p, 1).projective_eq(&two_p));
        // double(2P) == 4P
        assert!(
            st.double(&two_p).projective_eq(&expect_4p),
            "case {n}: double(2P) != 4P"
        );
    }
}

// 6. act_point (2-torsion translation) against the oracle

#[test]
fn act_point_oracle() {
    let doc = load(THETA_ARITH_VECTORS);
    let null = parse_coords(&doc["null_point"]);
    for case in doc["act_null_cases"].as_array().unwrap() {
        let i = case["i"].as_u64().unwrap() as usize;
        let j = case["j"].as_u64().unwrap() as usize;
        let expected = parse_coords(&case["act"]);
        let got = ThetaPointDim4::new(act_point(null.coords(), i, j));
        // act_point is an exact (non-projective) coordinate operation.
        for k in 0..16 {
            assert!(
                fp2_eq(&got.coords()[k], &expected.coords()[k]),
                "act_point(i={i}, j={j}) coord {k} mismatch"
            );
        }
    }
    // act_point(P, 0, 0) is the identity.
    let id = ThetaPointDim4::new(act_point(null.coords(), 0, 0));
    for k in 0..16 {
        assert!(fp2_eq(&id.coords()[k], &null.coords()[k]));
    }
}
