use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use std::sync::Arc;
use symclaw_core::ast::Expr;
use symclaw_core::differentiate::{differentiate, nth_derivative};
use symclaw_core::egraph::egraph_simplify;
use symclaw_core::integrate::integrate;
use symclaw_core::interner::Symbol;
use symclaw_core::latex::to_latex;
use symclaw_core::parser::parse;
use symclaw_core::series::{maclaurin, taylor};
use symclaw_core::simplify::simplify;
use symclaw_core::solve::solve;

// ── Helpers ──────────────────────────────────────────────────────

fn p(input: &str) -> Arc<Expr> {
    parse(input).unwrap_or_else(|e| panic!("parse failed for {input:?}: {e}"))
}

fn sym(name: &str) -> Symbol {
    Symbol::new(name)
}

// ── Parsing ──────────────────────────────────────────────────────

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

    let cases = [
        ("simple", "x + 1"),
        ("medium", "x^3 * sin(x) + 2*x^2 - cos(x/2)"),
        ("complex", "(x^4 + 3*x^3 - 2*x^2 + x - 1) / (x^2 + 1)"),
    ];

    for (name, input) in &cases {
        group.bench_with_input(BenchmarkId::new("parse", name), input, |b, input| {
            b.iter(|| parse(black_box(input)).unwrap());
        });
    }

    group.finish();
}

// ── Simplification ──────────────────────────────────────────────

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

    // Identity removal
    let identity_cases = [
        ("add_zero", "x + 0"),
        ("mul_one", "x * 1"),
        ("pow_one", "x^1"),
    ];
    for (name, input) in &identity_cases {
        let expr = p(input);
        group.bench_with_input(BenchmarkId::new("identity", name), &expr, |b, expr| {
            b.iter(|| simplify(black_box(expr)));
        });
    }

    // Constant folding
    {
        let expr = p("2 + 3 + 4 + 5");
        group.bench_function("constant_folding", |b| {
            b.iter(|| simplify(black_box(&expr)));
        });
    }

    // Like-term combination
    {
        let expr = p("3*x + 5*x + 2*x");
        group.bench_function("like_terms", |b| {
            b.iter(|| simplify(black_box(&expr)));
        });
    }

    // Complex via e-graph
    {
        let expr = p("(x^2 - 1) / (x - 1)");
        group.bench_function("egraph_complex", |b| {
            b.iter(|| egraph_simplify(black_box(&expr)));
        });
    }

    group.finish();
}

// ── Differentiation ─────────────────────────────────────────────

fn bench_differentiation(c: &mut Criterion) {
    let mut group = c.benchmark_group("differentiation");
    let x = sym("x");

    // Polynomial
    {
        let expr = p("x^5 + 3*x^3 - x");
        group.bench_function("polynomial", |b| {
            b.iter(|| differentiate(black_box(&expr), x));
        });
    }

    // Trigonometric
    {
        let expr = p("sin(x^2) * cos(x)");
        group.bench_function("trigonometric", |b| {
            b.iter(|| differentiate(black_box(&expr), x));
        });
    }

    // Chain rule depth
    {
        let expr = p("sin(cos(exp(x)))");
        group.bench_function("chain_rule_deep", |b| {
            b.iter(|| differentiate(black_box(&expr), x));
        });
    }

    // Higher-order: d^5/dx^5(x^6)
    {
        let expr = p("x^6");
        group.bench_function("higher_order_5th", |b| {
            b.iter(|| nth_derivative(black_box(&expr), x, 5));
        });
    }

    group.finish();
}

// ── Integration ─────────────────────────────────────────────────

fn bench_integration(c: &mut Criterion) {
    let mut group = c.benchmark_group("integration");
    let x = sym("x");

    // Power rule
    {
        let expr = p("x^4");
        group.bench_function("power_rule", |b| {
            b.iter(|| integrate(black_box(&expr), x));
        });
    }

    // Trig
    {
        let expr = p("sin(x) * cos(x)");
        group.bench_function("trig", |b| {
            b.iter(|| integrate(black_box(&expr), x));
        });
    }

    // By parts
    {
        let expr = p("x * exp(x)");
        group.bench_function("by_parts", |b| {
            b.iter(|| integrate(black_box(&expr), x));
        });
    }

    group.finish();
}

// ── Solving ─────────────────────────────────────────────────────

fn bench_solving(c: &mut Criterion) {
    let mut group = c.benchmark_group("solving");
    let x = sym("x");

    // Linear: 2*x + 6 = 0
    {
        let expr = p("2*x + 6");
        group.bench_function("linear", |b| {
            b.iter(|| solve(black_box(&expr), x));
        });
    }

    // Quadratic: x^2 - 5*x + 6 = 0
    {
        let expr = p("x^2 - 5*x + 6");
        group.bench_function("quadratic", |b| {
            b.iter(|| solve(black_box(&expr), x));
        });
    }

    // Cubic: x^3 - 6*x^2 + 11*x - 6 = 0
    {
        let expr = p("x^3 - 6*x^2 + 11*x - 6");
        group.bench_function("cubic", |b| {
            b.iter(|| solve(black_box(&expr), x));
        });
    }

    group.finish();
}

// ── Series ──────────────────────────────────────────────────────

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

    // Maclaurin sin(x) order 10
    {
        let expr = p("sin(x)");
        group.bench_function("maclaurin_sin_10", |b| {
            b.iter(|| maclaurin(black_box(&expr), x, 10));
        });
    }

    // Taylor exp(x) about 1, order 8
    {
        let expr = p("exp(x)");
        let point = Expr::from(1i64);
        group.bench_function("taylor_exp_about1_8", |b| {
            b.iter(|| taylor(black_box(&expr), x, black_box(&point), 8));
        });
    }

    // Maclaurin x * sin(x) order 6
    {
        let expr = p("x * sin(x)");
        group.bench_function("maclaurin_xsinx_6", |b| {
            b.iter(|| maclaurin(black_box(&expr), x, 6));
        });
    }

    group.finish();
}

// ── E-graph ─────────────────────────────────────────────────────

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

    // Simple
    {
        let expr = p("x * 1 + 0");
        group.bench_function("simple_identity", |b| {
            b.iter(|| egraph_simplify(black_box(&expr)));
        });
    }

    // Trig identity
    {
        let expr = p("sin(x)^2 + cos(x)^2");
        group.bench_function("trig_identity", |b| {
            b.iter(|| egraph_simplify(black_box(&expr)));
        });
    }

    // Distribution
    {
        let expr = p("x * (y + z)");
        group.bench_function("distribution", |b| {
            b.iter(|| egraph_simplify(black_box(&expr)));
        });
    }

    group.finish();
}

// ── LaTeX Rendering ─────────────────────────────────────────────

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

    // Simple
    {
        let expr = p("x^2 + 3*x - 7");
        group.bench_function("simple", |b| {
            b.iter(|| to_latex(black_box(&expr)));
        });
    }

    // Complex fraction
    {
        let expr = p("(x^4 + 3*x^3 - 2*x^2 + x - 1) / (x^2 + 1)");
        group.bench_function("complex_fraction", |b| {
            b.iter(|| to_latex(black_box(&expr)));
        });
    }

    group.finish();
}

// ── Criterion Harness ───────────────────────────────────────────

criterion_group!(
    benches,
    bench_parsing,
    bench_simplification,
    bench_differentiation,
    bench_integration,
    bench_solving,
    bench_series,
    bench_egraph,
    bench_latex,
);
criterion_main!(benches);