use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use num_rational::Rational64;
use std::sync::Arc;
use symclaw_core::ast::Expr;
use symclaw_core::interner::Symbol;

// ============================================================
// PARSING BENCHMARKS
// ============================================================

fn bench_parse(c: &mut Criterion) {
    let mut group = c.benchmark_group("parse");

    let exprs = [
        ("simple", "x + 1"),
        ("medium", "x^3 + 2*x^2 + 3*x + 4"),
        ("complex", "sin(x)^2 + cos(x)^2 + exp(x*y) + ln(x + 1)"),
        ("deep", "((((x+1)^2+2)^2+3)^2+4)^2"),
    ];

    for (name, expr) in &exprs {
        group.bench_with_input(BenchmarkId::new("expr", name), expr, |b, expr| {
            b.iter(|| symclaw_core::parser::parse(black_box(expr)).unwrap())
        });
    }
    group.finish();
}

// ============================================================
// SIMPLIFICATION BENCHMARKS
// ============================================================

fn bench_simplify(c: &mut Criterion) {
    let mut group = c.benchmark_group("simplify");

    let cases = [
        ("identity", "x + 0"),
        ("combine", "2*x + 3*x"),
        ("expand", "(x+1)*(x-1)"),
        ("trig", "sin(x)^2 + cos(x)^2"),
        ("nested", "(x+1)^2 - x^2 - 2*x - 1"),
    ];

    for (name, expr_str) in &cases {
        let expr = symclaw_core::parser::parse(expr_str).unwrap();
        group.bench_with_input(BenchmarkId::new("expr", name), &expr, |b, expr| {
            b.iter(|| symclaw_core::simplify::simplify(black_box(expr)))
        });
    }
    group.finish();
}

// ============================================================
// DIFFERENTIATION BENCHMARKS
// ============================================================

fn bench_differentiate(c: &mut Criterion) {
    let mut group = c.benchmark_group("differentiate");

    let x = Symbol::new("x");

    let cases = [
        ("polynomial", "x^5 + 3*x^3 + 2*x"),
        ("trig", "sin(x)*cos(x)"),
        ("chain", "sin(cos(exp(x)))"),
        ("product", "x^2*sin(x)*exp(x)"),
    ];

    for (name, expr_str) in &cases {
        let expr = symclaw_core::parser::parse(expr_str).unwrap();
        group.bench_with_input(BenchmarkId::new("expr", name), &expr, |b, expr| {
            b.iter(|| symclaw_core::differentiate::differentiate(black_box(expr), x))
        });
    }
    group.finish();
}

// ============================================================
// INTEGRATION BENCHMARKS
// ============================================================

fn bench_integrate(c: &mut Criterion) {
    let mut group = c.benchmark_group("integrate");

    let x = Symbol::new("x");

    let cases = [
        ("polynomial", "x^4 + 2*x^2 + 1"),
        ("trig", "sin(x)"),
        ("exp", "exp(x)"),
    ];

    for (name, expr_str) in &cases {
        let expr = symclaw_core::parser::parse(expr_str).unwrap();
        group.bench_with_input(BenchmarkId::new("expr", name), &expr, |b, expr| {
            b.iter(|| symclaw_core::integrate::integrate(black_box(expr), x))
        });
    }
    group.finish();
}

// ============================================================
// POLYNOMIAL GCD BENCHMARKS — THE KEY BENCHMARK
// ============================================================

fn bench_polynomial_gcd(c: &mut Criterion) {
    use symclaw_core::poly::Polynomial;

    let mut group = c.benchmark_group("polynomial_gcd");
    group.sample_size(10);

    // Symbolica-style benchmark: multivariate GCD in 7 variables
    // a = (1 + 3*x1 + 5*x2 + 7*x3 + 9*x4 + 11*x5 + 13*x6 + 15*x7)^3 - 1
    // b = (1 - 3*x1 - 5*x2 - 7*x3 + 9*x4 - 11*x5 - 13*x6 + 15*x7)^3 + 1
    // g = (1 + 3*x1 + 5*x2 + 7*x3 + 9*x4 + 11*x5 + 13*x6 - 15*x7)^3 + 3
    // Compute: gcd(a*g, b*g) — should equal g (up to scalar)

    let vars: Vec<String> = (1..=7).map(|i| format!("x{i}")).collect();

    fn build_poly(coeffs: &[i64], vars: &[String], deg: u32, add_const: i64) -> Polynomial {
        let mut base = Polynomial::constant(Rational64::from_integer(coeffs[0]), vars.to_vec());
        for (i, &c) in coeffs[1..].iter().enumerate() {
            let term = Polynomial::var(&vars[i], vars.to_vec()).scale(Rational64::from_integer(c));
            base = base.add(&term);
        }
        let mut result = base.pow(deg);
        if add_const != 0 {
            result = result.add(&Polynomial::constant(
                Rational64::from_integer(add_const),
                vars.to_vec(),
            ));
        }
        result
    }

    // Degree 3 (manageable for criterion timing)
    let a = build_poly(&[1, 3, 5, 7, 9, 11, 13, 15], &vars, 3, -1);
    let b = build_poly(&[1, -3, -5, -7, 9, -11, -13, 15], &vars, 3, 1);
    let g = build_poly(&[1, 3, 5, 7, 9, 11, 13, -15], &vars, 3, 3);
    let ag = a.mul(&g);
    let bg = b.mul(&g);

    group.bench_function("subresultant_prs_deg3_7var", |bench| {
        bench.iter(|| black_box(&ag).gcd(black_box(&bg)))
    });

    group.bench_function("modular_gcd_deg3_7var", |bench| {
        bench.iter(|| symclaw_core::modular_gcd::modular_gcd(black_box(&ag), black_box(&bg)))
    });

    // Univariate GCD: (1+x)^10 vs (1+x)^8 → gcd = (1+x)^8
    let uv = vec!["x".to_owned()];
    let p1 = {
        let base = Polynomial::constant(Rational64::from_integer(1), uv.clone())
            .add(&Polynomial::var("x", uv.clone()));
        base.pow(10)
    };
    let p2 = {
        let base = Polynomial::constant(Rational64::from_integer(1), uv.clone())
            .add(&Polynomial::var("x", uv.clone()));
        base.pow(8)
    };

    group.bench_function("univariate_gcd_deg10", |bench| {
        bench.iter(|| black_box(&p1).gcd(black_box(&p2)))
    });

    group.finish();
}

// ============================================================
// PACKED POLY BENCHMARKS
// ============================================================

fn bench_packed_poly(c: &mut Criterion) {
    use symclaw_core::packed_poly::PackedPoly;
    use symclaw_core::poly::Polynomial;

    let mut group = c.benchmark_group("packed_poly");

    let vars = vec!["x".to_owned(), "y".to_owned()];
    let mut poly = Polynomial::zero(vars.clone());
    for i in 0..50u32 {
        for j in 0..10u32 {
            let coeff = Rational64::from_integer((i as i64 + 1) * (j as i64 + 1));
            let term = Polynomial::constant(coeff, vars.clone())
                .mul(&Polynomial::var("x", vars.clone()).pow(i))
                .mul(&Polynomial::var("y", vars.clone()).pow(j));
            poly = poly.add(&term);
        }
    }

    group.bench_function("tree_to_packed", |b| {
        b.iter(|| PackedPoly::from_polynomial(black_box(&poly)))
    });

    let packed = PackedPoly::from_polynomial(&poly);
    group.bench_function("packed_to_tree", |b| {
        b.iter(|| black_box(&packed).to_polynomial())
    });

    group.bench_function("packed_add", |b| {
        b.iter(|| black_box(&packed).add(black_box(&packed)))
    });

    group.finish();
}

// ============================================================
// SIMD ARITHMETIC BENCHMARKS
// ============================================================

fn bench_simd_arith(c: &mut Criterion) {
    use symclaw_core::simd_arith::*;

    let mut group = c.benchmark_group("simd_arith");

    let n = 10_000usize;
    let modulus = 998_244_353u64;
    let a: Vec<u64> = (0..n as u64).map(|i| i % modulus).collect();
    let b: Vec<u64> = (0..n as u64).map(|i| (i * 7 + 3) % modulus).collect();
    let mut result = vec![0u64; n];

    group.bench_function("batch_mod_mul_10k", |bench| {
        bench.iter(|| batch_mod_mul(black_box(&a), black_box(&b), modulus, &mut result))
    });

    group.bench_function("batch_mod_add_10k", |bench| {
        bench.iter(|| batch_mod_add(black_box(&a), black_box(&b), modulus, &mut result))
    });

    let ctx = MontgomeryCtx::new(modulus);
    let a_mont: Vec<u64> = a.iter().map(|&x| ctx.to_mont(x)).collect();
    let b_mont: Vec<u64> = b.iter().map(|&x| ctx.to_mont(x)).collect();
    let mut r_mont = vec![0u64; n];

    group.bench_function("batch_montgomery_mul_10k", |bench| {
        bench.iter(|| ctx.batch_mont_mul(black_box(&a_mont), black_box(&b_mont), &mut r_mont))
    });

    let poly_a: Vec<u64> = (0..1000u64).map(|i| i % modulus).collect();
    let poly_b: Vec<u64> = (0..1000u64).map(|i| (i * 3 + 1) % modulus).collect();

    group.bench_function("poly_mul_mod_1000", |bench| {
        bench.iter(|| poly_mul_mod(black_box(&poly_a), black_box(&poly_b), modulus))
    });

    group.finish();
}

// ============================================================
// E-GRAPH BENCHMARKS
// ============================================================

fn bench_egraph(c: &mut Criterion) {
    let mut group = c.benchmark_group("egraph");

    let cases = [
        ("simple", "x + 0"),
        ("medium", "x * 1 + 0 * y"),
        ("complex", "(a + b) * (a + b) - a*a - 2*a*b - b*b"),
    ];

    for (name, expr_str) in &cases {
        let expr = symclaw_core::parser::parse(expr_str).unwrap();
        group.bench_with_input(BenchmarkId::new("simplify", name), &expr, |b, expr| {
            b.iter(|| symclaw_core::egraph::egraph_simplify(black_box(expr)))
        });
    }
    group.finish();
}

// ============================================================
// SERIES BENCHMARKS
// ============================================================

fn bench_series(c: &mut Criterion) {
    let mut group = c.benchmark_group("series");

    let x = Symbol::new("x");
    let zero = Arc::new(Expr::Num(Rational64::from_integer(0)));

    let expr_sin = symclaw_core::parser::parse("sin(x)").unwrap();
    for order in [5u32, 10, 20] {
        group.bench_with_input(
            BenchmarkId::new("taylor_sin", order),
            &order,
            |b, &order| {
                b.iter(|| {
                    symclaw_core::series::taylor(black_box(&expr_sin), x, &zero, black_box(order))
                })
            },
        );
    }

    let expr_exp = symclaw_core::parser::parse("exp(x)").unwrap();
    for order in [5u32, 10, 20] {
        group.bench_with_input(
            BenchmarkId::new("taylor_exp", order),
            &order,
            |b, &order| {
                b.iter(|| {
                    symclaw_core::series::taylor(black_box(&expr_exp), x, &zero, black_box(order))
                })
            },
        );
    }

    group.finish();
}

// ============================================================

criterion_group!(
    benches,
    bench_parse,
    bench_simplify,
    bench_differentiate,
    bench_integrate,
    bench_polynomial_gcd,
    bench_packed_poly,
    bench_simd_arith,
    bench_egraph,
    bench_series,
);
criterion_main!(benches);