zyga 0.5.1

use crate::common::{Proof, TrustedSetup};
use crate::debug_println;
use ark_bn254::{Fr, G1Affine, G1Projective, G2Affine, G2Projective};
use ark_ec::{AffineRepr, CurveGroup, Group};
use ark_ff::UniformRand;
use ark_serialize::CanonicalSerialize;
use ark_std::{rand::RngCore, Zero};
use std::collections::HashMap;
use std::ops::Neg;

use crate::{
    code_generation::generate_public_coefficients_file,
    common::{G1Point, G2Point},
    constraint::CompilationResult,
    dag::Expression,
    errors::ZkError,
    polynomial::vanishing_polynomial,
    proving_key::{ProvingKey, VerificationKey},
};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PairingProof {
    pub proof: Proof,
    pub trusted_setup: TrustedSetup,
    pub n_constraints: usize,
}

/// Generate trusted setup parameters
pub fn generate_trusted_setup<R: RngCore>(rng: &mut R) -> Result<TrustedSetup, ZkError> {
    // Generate random field elements
    let alpha = Fr::rand(rng);
    let beta = Fr::rand(rng);

    // Get generators
    let g1 = G1Projective::generator();
    let g2 = G2Projective::generator();

    // Compute h1 = g1^alpha and h2 = g2^beta
    let h1 = g1 * alpha;
    let h2 = g2 * beta;

    // Serialize points
    let h1_bytes = serialize_g1_point(&h1.into_affine())?;
    let h2_bytes = serialize_g2_point(&h2.into_affine())?;

    Ok(TrustedSetup {
        alpha: alpha.to_string(),
        beta: beta.to_string(),
        h1: h1_bytes,
        h2: h2_bytes,
    })
}

/// Create proving key and verification key from compiled constraints
pub fn create_proving_key(
    compilation_result: &CompilationResult,
    trusted_setup: &TrustedSetup,
    constraint_file_path: Option<&std::path::Path>,
    prefix: Option<&str>,
) -> Result<(ProvingKey, VerificationKey), ZkError> {
    let n_constraints = compilation_result.a_matrix.len();

    // Trusted setup and generators are not required during setup for codegen

    // Evaluate at a single point outside the constraint set
    // We use n_constraints + 1 as the evaluation point
    let eval_point = (n_constraints + 1) as f64;

    // Create Lagrange polynomials for constraint matrices
    use crate::polynomial::lagrange_interpolate;

    // Evaluate Lagrange polynomials at the evaluation point
    let a_poly_at_point = lagrange_interpolate(&compilation_result.a_matrix, eval_point);
    let b_poly_at_point = lagrange_interpolate(&compilation_result.b_matrix, eval_point);
    let c_poly_at_point = lagrange_interpolate(&compilation_result.c_matrix, eval_point);

    debug_println!("\n=== Lagrange Coefficients at x={} ===", eval_point);
    debug_println!(
        "A coefficients (first 5): {:?}",
        &a_poly_at_point[..5.min(a_poly_at_point.len())]
    );
    debug_println!(
        "B coefficients (first 5): {:?}",
        &b_poly_at_point[..5.min(b_poly_at_point.len())]
    );

    // Use HashMap to accumulate public coefficients by symbol name to handle duplicates
    let mut public_coeffs_a: HashMap<String, f64> = HashMap::new();
    let mut public_coeffs_b: HashMap<String, f64> = HashMap::new();
    let mut public_coeffs_c: HashMap<String, f64> = HashMap::new();

    debug_println!("\n=== Computing Witness Polynomials ===");
    debug_println!("First 10 witness values and their contributions:");

    for (i, &witness_id) in compilation_result.witness_ids.iter().enumerate() {
        if i >= a_poly_at_point.len() {
            break;
        }

        let coeff_a = a_poly_at_point[i];
        let coeff_b = b_poly_at_point[i];
        let coeff_c = c_poly_at_point[i];

        let witness_expr = compilation_result.dag.get(witness_id);

        if let Expression::Deferred(name) = witness_expr {
                // Public symbolic value - coefficient only (NOT multiplied by actual value)
                // The actual multiplication happens at verification time

                // Accumulate coefficients by symbol name to handle duplicates properly
                // This matches Python's behavior where each unique symbol contributes once
                *public_coeffs_a.entry(name.clone()).or_insert(0.0) += coeff_a;
                *public_coeffs_b.entry(name.clone()).or_insert(0.0) += coeff_b;
                *public_coeffs_c.entry(name.clone()).or_insert(0.0) += coeff_c;
        } else {
            // Only deferred (public) symbols are needed for codegen
        }
    }

    // Check if all B coefficients are zero (would result in trivial pairing)
    let b_all_zero = public_coeffs_b.values().all(|&coeff| coeff.abs() < 1e-10);

    if b_all_zero {
        debug_println!("WARNING: All public B coefficients are zero.");
        debug_println!("This indicates the constraint system needs public/deferred values in B position.");
        debug_println!("The dummy constraint injection in compile_constraints should have handled this.");
    }

    // Public coefficients accumulated; generate codegen file

    // Generate public_coefficients.rs file for on-chain verification
    generate_public_coefficients_file(
        &public_coeffs_a,
        &public_coeffs_b,
        &public_coeffs_c,
        constraint_file_path,
        prefix,
    )?;

    // Convert env_dict to integer values (inputs are always integers)
    let mut env_dict_i64 = HashMap::new();
    for (key, value) in &compilation_result.env_dict {
        if let Some(val) = value.as_f64() {
            env_dict_i64.insert(key.clone(), val as i64);  // Convert to integer
        }
    }

    // Create proving key with pre-computed Lagrange evaluations
    let proving_key = ProvingKey {
        a_matrix: compilation_result.a_matrix.clone(),
        b_matrix: compilation_result.b_matrix.clone(),
        c_matrix: compilation_result.c_matrix.clone(),
        lagrange_a: a_poly_at_point,
        lagrange_b: b_poly_at_point,
        lagrange_c: c_poly_at_point,
        witness_dag: compilation_result.dag.clone(),
        witness_ids: compilation_result.witness_ids.clone(),
        witness_names: compilation_result.witnesses.clone(),
        public_variables: compilation_result.public_variables.clone(),
        trusted_setup: trusted_setup.clone(),
        num_constraints: n_constraints,
        num_variables: compilation_result.witnesses.len(),
        evaluation_point: eval_point,
        env_dict: env_dict_i64,
    };

    // Create verification key with static elements
    let g1 = G1Projective::generator();
    let g2 = G2Projective::generator();
    let g1_bytes = serialize_g1_point(&g1.into_affine())?;
    let g2_bytes = serialize_g2_point(&g2.into_affine())?;

    let verification_key = VerificationKey {
        g1: G1Point::from_vec(g1_bytes).map_err(ZkError::ComputationError)?,
        g2: G2Point::from_vec(g2_bytes).map_err(ZkError::ComputationError)?,
        h1: G1Point::from_vec(trusted_setup.h1.clone()).map_err(ZkError::ComputationError)?,
        h2: G2Point::from_vec(trusted_setup.h2.clone()).map_err(ZkError::ComputationError)?,
        num_constraints: n_constraints,
    };

    Ok((proving_key, verification_key))
}

/// Generate a proof using proving key and witness values
pub fn generate_proof(
    proving_key: &ProvingKey,
    expanded_witness: &HashMap<String, f64>,
    force_proof: bool,
) -> Result<PairingProof, ZkError> {
    let n_constraints = proving_key.num_constraints;
    let eval_point = proving_key.evaluation_point;

    // Parse trusted setup
    let h1 = deserialize_g1_point(&proving_key.trusted_setup.h1)?;

    // Get standard generators for public parts
    let g1 = G1Projective::generator();
    let g2 = G2Projective::generator();

    // Use pre-computed Lagrange evaluations from proving key
    let a_poly_at_point = &proving_key.lagrange_a;
    let b_poly_at_point = &proving_key.lagrange_b;
    let c_poly_at_point = &proving_key.lagrange_c;

    debug_println!("\n=== Lagrange Coefficients at x={} ===", eval_point);
    debug_println!(
        "A coefficients (first 5): {:?}",
        &a_poly_at_point[..5.min(a_poly_at_point.len())]
    );
    debug_println!(
        "B coefficients (first 5): {:?}",
        &b_poly_at_point[..5.min(b_poly_at_point.len())]
    );

    // Now separate private and public contributions by going through witness values
    let mut a_private = 0.0;
    let mut b_private = 0.0;
    let mut c_private = 0.0;

    // Use HashMap to accumulate public coefficients by symbol name to handle duplicates
    let mut public_coeffs_a: HashMap<String, f64> = HashMap::new();
    let mut public_coeffs_b: HashMap<String, f64> = HashMap::new();
    let mut public_coeffs_c: HashMap<String, f64> = HashMap::new();

    debug_println!("\n=== Computing Witness Polynomials ===");
    debug_println!("First 10 witness values and their contributions:");

    for (i, &witness_id) in proving_key.witness_ids.iter().enumerate() {
        if i >= a_poly_at_point.len() {
            break;
        }

        let coeff_a = a_poly_at_point[i];
        let coeff_b = b_poly_at_point[i];
        let coeff_c = c_poly_at_point[i];

        let witness_expr = proving_key.witness_dag.get(witness_id);

        if i < 20 {
            debug_println!("  [{:2}] witness: {:?}", i, witness_expr);
            debug_println!(
                ", coeff_A: {:.2}, coeff_B: {:.2}, coeff_C: {:.2}",
                coeff_a,
                coeff_b,
                coeff_c
            );
        }

        match witness_expr {
            Expression::Private(name) => {
                // Private witness value - get from expanded_witness
                let val = expanded_witness.get(name)
                    .ok_or_else(|| ZkError::InvalidParameters)?;
                a_private += coeff_a * val;
                b_private += coeff_b * val;
                c_private += coeff_c * val;
            }
            Expression::Deferred(name) => {
                // Public symbolic value - coefficient only (NOT multiplied by actual value)
                // The actual multiplication happens at verification time

                if i < 10 {
                    debug_println!("    Deferred {} (coefficient only for verification)", name);
                }

                // Accumulate coefficients by symbol name to handle duplicates properly
                // This matches Python's behavior where each unique symbol contributes once
                *public_coeffs_a.entry(name.clone()).or_insert(0.0) += coeff_a;
                *public_coeffs_b.entry(name.clone()).or_insert(0.0) += coeff_b;
                *public_coeffs_c.entry(name.clone()).or_insert(0.0) += coeff_c;
            }
            Expression::Public(name) => {
                // Public witness value - get from expanded_witness
                let val = expanded_witness.get(name)
                    .ok_or_else(|| ZkError::InvalidParameters)?;
                a_private += coeff_a * val;
                b_private += coeff_b * val;
                c_private += coeff_c * val;
            }
            Expression::Constant(val) => {
                // Literal constants from constraints
                a_private += coeff_a * val.0;
                b_private += coeff_b * val.0;
                c_private += coeff_c * val.0;
            }
            _ => {
                // Complex expression - try to evaluate if possible
                if proving_key.witness_dag.can_evaluate(witness_id) {
                    let val = proving_key.witness_dag.evaluate(witness_id);
                    a_private += coeff_a * val;
                    b_private += coeff_b * val;
                    c_private += coeff_c * val;
                } else {
                    // Expression contains deferred values - this is expected and normal
                    // These will be handled by the public coefficients at verification time
                }
            }
        }
    }

    // Compute public contributions using actual witness values
    // Create an updated env_dict with witness values (as integers)
    let mut env_dict = proving_key.env_dict.clone();

    // Special case: the deferred constant "1" always has value 1
    env_dict.insert("1".to_string(), 1);

    // Update env_dict with expanded witness values (convert to integers)
    // This includes both the original values and expanded array elements
    for (key, &value) in expanded_witness {
        env_dict.insert(key.clone(), value as i64);
    }

    debug_println!("\n=== Updated env_dict with witness values ===");
    debug_println!("env_dict now has {} entries", env_dict.len());
    if env_dict.contains_key("a[0]") {
        debug_println!("Found a[0] in env_dict: {}", env_dict["a[0]"]);
    } else {
        debug_println!("WARNING: a[0] NOT found in env_dict");
        debug_println!("Keys in env_dict: {:?}", env_dict.keys().collect::<Vec<_>>());
    }

    // Compute public contributions using i64-wrapping arithmetic to match on-chain codegen
    let mut a2_i: i64 = 0;
    let mut b2_i: i64 = 0;
    let mut c2_i: i64 = 0;

    for (symbol, coeff_a_f) in &public_coeffs_a {
        let coeff_a = *coeff_a_f as i64; // mirror codegen rounding
        let actual_value = env_dict.get(symbol).copied().unwrap_or(0) as i64;
        a2_i = a2_i.wrapping_add(coeff_a.wrapping_mul(actual_value));
    }
    for (symbol, coeff_b_f) in &public_coeffs_b {
        let coeff_b = *coeff_b_f as i64;
        let actual_value = env_dict.get(symbol).copied().unwrap_or(0) as i64;
        b2_i = b2_i.wrapping_add(coeff_b.wrapping_mul(actual_value));
    }
    for (symbol, coeff_c_f) in &public_coeffs_c {
        let coeff_c = *coeff_c_f as i64;
        let actual_value = env_dict.get(symbol).copied().unwrap_or(0) as i64;
        c2_i = c2_i.wrapping_add(coeff_c.wrapping_mul(actual_value));
    }

    // Debug: Show accumulated coefficients by symbol (floats still useful for tracing)
    debug_println!("\n=== Accumulated Public Coefficients ===");
    for (symbol, coeff) in &public_coeffs_a {
        if coeff.abs() > 1e-10 {
            debug_println!("  {}: A={:.2}", symbol, coeff);
        }
    }

    debug_println!("\n=== Private vs Public Split Debug ===");
    debug_println!("  a_private = {} (should match Python a1)", a_private);
    debug_println!("  a_public  (i64) = {}", a2_i);
    debug_println!("  b_private = {} (should match Python b1)", b_private);
    debug_println!("  b_public  (i64) = {}", b2_i);
    debug_println!("  c_private = {} (should match Python c1)", c_private);
    debug_println!("  c_public  (i64) = {}", c2_i);

    // Convert to integers (prototype path) then into Fr for safe field arithmetic
    let a1_i = a_private as i64;
    let b1_i = b_private as i64;
    let c1_i = c_private as i64;

    // Map i64 -> Fr (signed mapping)
    #[inline]
    fn fr_from_i64(x: i64) -> Fr {
        if x >= 0 {
            Fr::from(x as u64)
        } else {
            -Fr::from((-x) as u64)
        }
    }

    let a1 = fr_from_i64(a1_i);
    let a2 = fr_from_i64(a2_i);
    let b1 = fr_from_i64(b1_i);
    let b2 = fr_from_i64(b2_i);
    let c1 = fr_from_i64(c1_i);
    let c2 = fr_from_i64(c2_i);

    // First, check constraint satisfaction with actual witness values
    debug_println!("\n=== Checking Constraint Satisfaction ===");
    let mut all_satisfied = true;
    for (i, ((a_row, b_row), c_row)) in proving_key
        .a_matrix
        .iter()
        .zip(proving_key.b_matrix.iter())
        .zip(proving_key.c_matrix.iter())
        .enumerate()
    {
        let mut a_val = 0.0;
        let mut b_val = 0.0;
        let mut c_val = 0.0;

        for (j, &witness_id) in proving_key.witness_ids.iter().enumerate() {
            if j >= a_row.len() {
                break;
            }

            let witness_val = match proving_key.witness_dag.get(witness_id) {
                Expression::Private(name) | Expression::Public(name) => {
                    // Look up value in env_dict
                    env_dict
                        .get(name)
                        .copied()
                        .ok_or_else(|| {
                            debug_println!("ERROR: Variable {} not found in witness", name);
                            ZkError::InvalidParameters
                        })? as f64
                }
                Expression::Constant(val) => val.0,
                Expression::Deferred(name) => {
                    // Look up deferred value in env_dict
                    env_dict
                        .get(name)
                        .copied()
                        .ok_or_else(|| {
                            debug_println!("ERROR: Deferred variable {} not found in witness", name);
                            ZkError::InvalidParameters
                        })? as f64
                }
                _ => {
                    // Try to evaluate complex expressions using env_dict
                    let mut env = HashMap::new();
                    for (key, val) in &env_dict {
                        env.insert(key.clone(), *val as f64);
                    }
                    debug_println!("Evaluating complex expression id {} with env containing {} keys", witness_id, env.len());
                    debug_println!("Expression is: {:?}", proving_key.witness_dag.get(witness_id));
                    proving_key
                        .witness_dag
                        .evaluate_with_env(witness_id, &env)
                        .map_err(|e| {
                            debug_println!("ERROR evaluating expression: {}", e);
                            debug_println!("Available keys in env: {:?}", env.keys().take(10).collect::<Vec<_>>());
                            ZkError::InvalidParameters
                        })?
                }
            };

            a_val += a_row[j] * witness_val;
            b_val += b_row[j] * witness_val;
            c_val += c_row[j] * witness_val;
        }

        let constraint_val: f64 = a_val * b_val - c_val;
        if constraint_val.abs() > 1e-6 {
            debug_println!(
                "  Constraint {}: A*B - C = {} * {} - {} = {} (NOT SATISFIED)",
                i,
                a_val,
                b_val,
                c_val,
                constraint_val
            );
            all_satisfied = false;
        }
    }

    if all_satisfied {
        debug_println!("  All {} constraints satisfied!", n_constraints);
    } else if force_proof {
        debug_println!("  WARNING: Some constraints not satisfied! Forcing proof generation due to --force flag.");
    } else {
        debug_println!("  WARNING: Some constraints not satisfied!");
        return Err(ZkError::InvalidParameters);
    }

    // Compute quotient polynomial H
    // H = (wA * wB - wC) / Z_n at eval_point
    // For debug: compute totals in f64 using the integer public parts
    let a_public_f = a2_i as f64;
    let b_public_f = b2_i as f64;
    let c_public_f = c2_i as f64;
    let wa_total = a_private + a_public_f;
    let wb_total = b_private + b_public_f;
    let wc_total = c_private + c_public_f;
    let p_val = wa_total * wb_total - wc_total;
    let z_val = vanishing_polynomial(n_constraints, eval_point);

    debug_println!("\n=== Polynomial Evaluation at x={} ===", eval_point);
    debug_println!(
        "  wA(x) = {} (private: {}, public: {})",
        wa_total,
        a_private,
        a_public_f
    );
    debug_println!(
        "  wB(x) = {} (private: {}, public: {})",
        wb_total,
        b_private,
        b_public_f
    );
    debug_println!(
        "  wC(x) = {} (private: {}, public: {})",
        wc_total,
        c_private,
        c_public_f
    );
    debug_println!("  P(x) = wA*wB - wC = {}", p_val);
    debug_println!("  Z(x) = {}", z_val);

    // OPTION A: Compute HZ = P directly in Fr (no division by Z)
    // This avoids the rounding-to-zero bug when Z is astronomical

    // Pure private term (goes to h1)
    let p_priv = a1 * b1 - c1; // a_priv * b_priv - c_priv (Fr)
                               // Mixed terms (both go to h1 in our routing)
    let p_mixed = a1 * b2 + a2 * b1; // a_priv * b_pub + a_pub * b_priv (Fr)
                                     // Pure public term (goes to g1)
    let p_pub = a2 * b2 - c2; // a_pub * b_pub - c_pub (Fr)

    debug_println!(
        "  P_priv = a1*b1 - c1 = {:?}*{:?} - {:?} = {:?}",
        a1,
        b1,
        c1,
        p_priv
    );
    debug_println!(
        "  P_mixed = a1*b2 + a2*b1 = {:?}*{:?} + {:?}*{:?} = {:?}",
        a1,
        b2,
        a2,
        b1,
        p_mixed
    );
    debug_println!(
        "  P_pub = a2*b2 - c2 = {:?}*{:?} - {:?} = {:?}",
        a2,
        b2,
        c2,
        p_pub
    );

    // HZ = P (no division needed!)
    // HZ_private gets both pure private and mixed terms (our routing: mixed → private)
    let hz_private_val = p_priv + p_mixed; // Fr
    let hz_public_val = p_pub; // Fr

    debug_println!(
        "  HZ_private = P_priv + P_mixed = {:?} + {:?} = {:?}",
        p_priv,
        p_mixed,
        hz_private_val
    );
    debug_println!("  HZ_public = P_pub = {:?}", hz_public_val);

    // Compute proof elements with ONE-SIDED PRIVATE ENCODING
    // All private mass goes to G1 only, no private in G2

    // A_curve = h1^a1 in G1 (private part) - unchanged
    let a_curve = scalar_mult_g1_fr(&h1, &a1)?;

    // B_curve = NO PRIVATE PART IN G2 (set to identity/zero)
    // We'll only use B_public = g2^b2
    // This removes the problematic e(h1, h2) cross term

    // C_curve = h1^c1 in G1 (private part) - unchanged initially
    let c_curve = scalar_mult_g1_fr(&h1, &c1)?;

    // Public parts (to be computed/verified on-chain with standard generators):
    // Pre-compute g1^a2 for on-chain verification (public part)
    let g1_a2 = scalar_mult_g1_fr(&g1, &a2)?;

    // Pre-compute g2^b2 for on-chain verification (public part, G2 scalar mul not available on-chain)
    // This is now our ONLY B component (no private B)
    let g2_b2 = scalar_mult_g2_fr(&g2, &b2)?;

    // Pre-compute g1^c2 for on-chain verification (public part)
    let g1_c2 = scalar_mult_g1_fr(&g1, &c2)?;

    // Compute HZ curve points - we already have hz_private_val and hz_public_val from P
    let h1_hz_private = scalar_mult_g1_fr(&h1, &hz_private_val)?; // Private HZ on h1
    let g1_hz_public = scalar_mult_g1_fr(&g1, &hz_public_val)?; // Public HZ on g1

    // ONE-SIDED ENCODING: Remove B_private entirely
    // In standard ZYGA: e(A_tot, B_tot) = e(C_tot, g2) * e(HZ_tot, g2)
    // With one-sided: e(A_tot, B_pub) = e(C'_tot, g2) * e(HZ'_tot, g2)
    // Where C' is adjusted to compensate for missing B_priv term

    // CRITICAL FIX: All compensation must be on the h1 (private) side
    // because mixed terms (a_pub * b_priv) are routed to the private bucket
    // Compute delta_priv = (a_priv + a_pub) * b_priv
    let delta_priv = (a1 + a2) * b1; // Total compensation on private side (Fr)

    debug_println!("\n=== ONE-SIDED ENCODING COMPENSATION ===");
    debug_println!("  a_priv (a1) = {}", a1);
    debug_println!("  a_pub (a2) = {}", a2);
    debug_println!("  b_priv (b1) = {}", b1);
    debug_println!("  b_pub (b2) = {}", b2);
    debug_println!("  c_priv (c1) = {}", c1);
    debug_println!("  c_pub (c2) = {}", c2);
    debug_println!(
        "  Delta_priv = (a_priv + a_pub) * b_priv = ({} + {}) * {} = {}",
        a1,
        a2,
        b1,
        delta_priv
    );
    debug_println!("  Original C scalar (c1): {}", c1);
    debug_println!("  Original HZ_private scalar: {}", hz_private_val);
    debug_println!("  Original HZ_public scalar: {}", hz_public_val);
    debug_println!(
        "  C' will have private part: {} - {} = {}",
        c1,
        delta_priv,
        c1 - delta_priv
    );

    // Sanity checks for exponent matching
    debug_println!("\n=== EXPONENT MATCHING SANITY CHECKS (WITH HZ=P) ===");

    // After compensation, we need:
    // LHS private channel: a_priv * b_pub (from e(A_total, B_pub))
    let lhs_private = a1 * b2;

    // RHS private channel: (C'_priv + HZ_priv) where C'_priv = c1 - delta_priv
    let c_prime_priv = c1 - delta_priv;
    let rhs_private = c_prime_priv + hz_private_val;

    debug_println!("  Private channel (h1):");
    debug_println!(
        "    LHS = a_priv * b_pub = {} * {} = {}",
        a1,
        b2,
        lhs_private
    );
    debug_println!(
        "    C'_priv = c1 - delta = {} - {} = {}",
        c1,
        delta_priv,
        c_prime_priv
    );
    debug_println!(
        "    RHS = C'_priv + HZ_priv = {} + {} = {}",
        c_prime_priv,
        hz_private_val,
        rhs_private
    );
    debug_println!(
        "    Match? {}",
        if lhs_private == rhs_private {
            "YES ✓"
        } else {
            "NO ✗"
        }
    );

    // LHS public channel: a_pub * b_pub (from e(A_total, B_pub))
    let lhs_public = a2 * b2;

    // RHS public channel: C_pub + HZ_pub (unchanged)
    let rhs_public = c2 + hz_public_val;

    debug_println!("  Public channel (g1):");
    debug_println!("    LHS = a_pub * b_pub = {} * {} = {}", a2, b2, lhs_public);
    debug_println!(
        "    RHS = C_pub + HZ_pub = {} + {} = {}",
        c2,
        hz_public_val,
        rhs_public
    );
    debug_println!(
        "    Match? {}",
        if lhs_public == rhs_public {
            "YES ✓"
        } else {
            "NO ✗"
        }
    );

    if lhs_private != rhs_private || lhs_public != rhs_public {
        debug_println!("  WARNING: Exponent mismatch detected! Pairing will fail.");
        debug_println!("  Gap on h1: {}", lhs_private - rhs_private);
        debug_println!("  Gap on g1: {}", lhs_public - rhs_public);
    }

    // Compute compensation point (entirely on h1)
    let delta_point = scalar_mult_g1_fr(&h1, &delta_priv)?;

    // Adjust C by subtracting delta from the private side
    // C'_private = C_private - delta_priv (on h1)
    // C'_public = C_public (unchanged, on g1)
    // HZ'_private = HZ_private (unchanged, on h1)
    // HZ'_public = HZ_public (unchanged, on g1)

    // For C: combine c_curve with negative delta_point (both on h1)
    // This gives us C'_private = C_private - delta_priv
    let c_private_adjusted = add_g1_points(&c_curve, &delta_point.neg())?;

    // HZ remains unchanged on both private and public sides
    let hz_private_final = h1_hz_private;
    let hz_public_final = g1_hz_public;

    // Serialize points
    let a_curve_bytes = serialize_g1_point(&a_curve)?;
    // NO b_curve for private - we only have public B
    let b_pub_bytes = serialize_g2_point(&g2_b2)?; // This is our only B component
    let c_private_adjusted_bytes = serialize_g1_point(&c_private_adjusted)?; // C'_private (adjusted)
    let g1_a2_bytes = serialize_g1_point(&g1_a2)?;
    let g1_c2_bytes = serialize_g1_point(&g1_c2)?;

    // Debug: show whether g1_a2/g1_c2 are infinity or non-zero
    debug_println!(
        "g1_a2 is {}",
        if g1_a2_bytes.iter().all(|&b| b == 0) { "INF (zero)" } else { "NON-ZERO" }
    );
    debug_println!(
        "g1_c2 is {}",
        if g1_c2_bytes.iter().all(|&b| b == 0) { "INF (zero)" } else { "NON-ZERO" }
    );

    // Combine HZ private and public to get total HZ
    // HZ_total = HZ_private + HZ_public
    let hz_total = add_g1_points(&hz_private_final, &hz_public_final)?;
    let hz_total_bytes = serialize_g1_point(&hz_total)?;

    // Create single proof element with one-sided encoding
    let proof = Proof {
        a_curve: G1Point::from_vec(a_curve_bytes).map_err(ZkError::ComputationError)?, // A_private
        g2_b2: G2Point::from_vec(b_pub_bytes).map_err(ZkError::ComputationError)?, // B_public
        c_curve: G1Point::from_vec(c_private_adjusted_bytes).map_err(ZkError::ComputationError)?, // C'_private
        g1_a2: G1Point::from_vec(g1_a2_bytes).map_err(ZkError::ComputationError)?, // A_public
        g1_c2: G1Point::from_vec(g1_c2_bytes).map_err(ZkError::ComputationError)?, // C_public
        g1_hz: G1Point::from_vec(hz_total_bytes).map_err(ZkError::ComputationError)?, // HZ_total
    };

    Ok(PairingProof { proof, trusted_setup: proving_key.trusted_setup.clone(), n_constraints })
}

/// Scalar multiplication in G1 with Fr scalar
fn scalar_mult_g1_fr(base: &G1Projective, scalar: &Fr) -> Result<G1Affine, ZkError> {
    if scalar.is_zero() {
        return Ok(G1Affine::zero());
    }
    let result = *base * *scalar;
    Ok(result.into_affine())
}

/// Scalar multiplication in G2 with Fr scalar
fn scalar_mult_g2_fr(base: &G2Projective, scalar: &Fr) -> Result<G2Affine, ZkError> {
    if scalar.is_zero() {
        return Ok(G2Affine::zero());
    }
    let result = *base * *scalar;
    Ok(result.into_affine())
}

/// Add two G1 points
fn add_g1_points(p1: &G1Affine, p2: &G1Affine) -> Result<G1Affine, ZkError> {
    let result = (*p1 + *p2).into_affine();
    Ok(result)
}

/// Serialize G1 point to Solana-compatible format (64 bytes)
fn serialize_g1_point(point: &G1Affine) -> Result<Vec<u8>, ZkError> {
    // Handle point at infinity
    if point.is_zero() {
        return Ok(vec![0u8; 64]);
    }

    let mut bytes = Vec::new();

    // Get x and y coordinates
    let x = *point.x().unwrap();
    let y = *point.y().unwrap();

    // Serialize x coordinate (32 bytes, little-endian)
    let mut x_bytes = vec![0u8; 32];
    x.serialize_uncompressed(&mut x_bytes[..])
        .map_err(|e| ZkError::ComputationError(format!("G1 x serialization failed: {:?}", e)))?;

    // Serialize y coordinate (32 bytes, little-endian)
    let mut y_bytes = vec![0u8; 32];
    y.serialize_uncompressed(&mut y_bytes[..])
        .map_err(|e| ZkError::ComputationError(format!("G1 y serialization failed: {:?}", e)))?;

    // Convert to big-endian for Solana
    x_bytes.reverse();
    y_bytes.reverse();

    bytes.extend_from_slice(&x_bytes);
    bytes.extend_from_slice(&y_bytes);

    Ok(bytes)
}

/// Serialize G2 point to Solana-compatible format (128 bytes)
/// Format follows EIP-197: [be(x1), be(x0), be(y1), be(y0)]
fn serialize_g2_point(point: &G2Affine) -> Result<Vec<u8>, ZkError> {
    let mut bytes = Vec::with_capacity(128);

    // Get x and y coordinates (each is an Fp2 element)
    let (x, y) = if point.is_zero() {
        // Point at infinity - return zeros
        bytes.resize(128, 0);
        return Ok(bytes);
    } else {
        (*point.x().unwrap(), *point.y().unwrap())
    };

    // G2 uses Fp2, so each coordinate has two components (c0, c1)
    // arkworks serializes as [le(x.c0), le(x.c1), le(y.c0), le(y.c1)]

    // Serialize x (64 bytes total)
    let mut x_bytes = [0u8; 64];
    x.serialize_uncompressed(&mut x_bytes[..])
        .map_err(|e| ZkError::ComputationError(format!("G2 x serialization failed: {:?}", e)))?;

    // Serialize y (64 bytes total)
    let mut y_bytes = [0u8; 64];
    y.serialize_uncompressed(&mut y_bytes[..])
        .map_err(|e| ZkError::ComputationError(format!("G2 y serialization failed: {:?}", e)))?;

    // arkworks gives us [le(x.c0), le(x.c1), le(y.c0), le(y.c1)]
    // We need EIP-197 format: [be(x.c1), be(x.c0), be(y.c1), be(y.c0)]

    // x.c1 in big-endian (bytes 32-64 of x_bytes, reversed)
    let mut x_c1_be = x_bytes[32..64].to_vec();
    x_c1_be.reverse();
    bytes.extend_from_slice(&x_c1_be);

    // x.c0 in big-endian (bytes 0-32 of x_bytes, reversed)
    let mut x_c0_be = x_bytes[0..32].to_vec();
    x_c0_be.reverse();
    bytes.extend_from_slice(&x_c0_be);

    // y.c1 in big-endian (bytes 32-64 of y_bytes, reversed)
    let mut y_c1_be = y_bytes[32..64].to_vec();
    y_c1_be.reverse();
    bytes.extend_from_slice(&y_c1_be);

    // y.c0 in big-endian (bytes 0-32 of y_bytes, reversed)
    let mut y_c0_be = y_bytes[0..32].to_vec();
    y_c0_be.reverse();
    bytes.extend_from_slice(&y_c0_be);

    Ok(bytes)
}

/// Deserialize G1 point from bytes
fn deserialize_g1_point(bytes: &[u8]) -> Result<G1Projective, ZkError> {
    if bytes.len() != 64 {
        return Err(ZkError::InvalidParameters);
    }

    // Check if all zeros (point at infinity)
    if bytes.iter().all(|&b| b == 0) {
        return Ok(G1Projective::zero());
    }

    // Parse x and y coordinates (big-endian format)
    let mut x_bytes = bytes[0..32].to_vec();
    let mut y_bytes = bytes[32..64].to_vec();

    // Convert from big-endian to little-endian
    x_bytes.reverse();
    y_bytes.reverse();

    // Deserialize field elements
    use ark_bn254::Fq;
    use ark_serialize::CanonicalDeserialize;

    let x = Fq::deserialize_uncompressed(&x_bytes[..])
        .map_err(|e| ZkError::ComputationError(format!("Failed to deserialize x: {:?}", e)))?;
    let y = Fq::deserialize_uncompressed(&y_bytes[..])
        .map_err(|e| ZkError::ComputationError(format!("Failed to deserialize y: {:?}", e)))?;

    // Create affine point and convert to projective
    let affine = G1Affine::new_unchecked(x, y);
    if !affine.is_on_curve() {
        return Err(ZkError::ComputationError("Point not on curve".to_string()));
    }

    Ok(affine.into())
}