//! Property-based tests for symclaw-core using proptest.
//!
//! Validates invariants of the symbolic math engine: simplification
//! preserves semantics, idempotency, parse round-trips, and more.

use std::collections::HashMap;
use std::sync::Arc;

use proptest::prelude::*;
use symclaw_core::ast::{Expr, FuncId};
use symclaw_core::cache::simplify_cached;
use symclaw_core::differentiate::differentiate;
use symclaw_core::eval::eval;
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::taylor;
use symclaw_core::simplify::simplify;
use symclaw_core::solve::solve;

// ── Configuration ────────────────────────────────────────────────

const PROPTEST_CASES: u32 = 500;

// ── Arbitrary Expr Generators ────────────────────────────────────

/// Generate a random expression tree up to `max_depth`.
/// Uses symbols x, y, z, constants 1–20, and operations
/// Add, Mul, Pow (exponents 1–4), Neg, Sin, Cos, Exp.
fn arb_expr(max_depth: u32) -> impl Strategy<Value = Arc<Expr>> {
    let leaf = prop_oneof![
        (1..=20i64).prop_map(Expr::int),
        Just("x").prop_map(Expr::sym),
        Just("y").prop_map(Expr::sym),
        Just("z").prop_map(Expr::sym),
    ];

    leaf.prop_recursive(max_depth, 64, 4, |inner| {
        prop_oneof![
            // Add: two children
            (inner.clone(), inner.clone()).prop_map(|(a, b)| Expr::add(vec![a, b])),
            // Mul: two children
            (inner.clone(), inner.clone()).prop_map(|(a, b)| Expr::mul(vec![a, b])),
            // Pow with small exponent
            (inner.clone(), 1..=4u32)
                .prop_map(|(base, exp)| Expr::pow(base, Expr::int(exp as i64))),
            // Neg
            inner.clone().prop_map(Expr::neg),
            // Sin
            inner.clone().prop_map(|e| Expr::func(FuncId::Sin, vec![e])),
            // Cos
            inner.clone().prop_map(|e| Expr::func(FuncId::Cos, vec![e])),
            // Exp
            inner.clone().prop_map(|e| Expr::func(FuncId::Exp, vec![e])),
        ]
    })
}

/// Generate a random polynomial (Add, Mul, Pow with integer exponents, Num, Sym only).
fn arb_polynomial(max_degree: u32) -> impl Strategy<Value = Arc<Expr>> {
    let leaf = prop_oneof![
        (1..=10i64).prop_map(Expr::int),
        Just("x").prop_map(Expr::sym),
        Just("y").prop_map(Expr::sym),
    ];

    leaf.prop_recursive(max_degree, 32, 3, |inner| {
        prop_oneof![
            (inner.clone(), inner.clone()).prop_map(|(a, b)| Expr::add(vec![a, b])),
            (inner.clone(), inner.clone()).prop_map(|(a, b)| Expr::mul(vec![a, b])),
            (inner.clone(), 1..=4u32)
                .prop_map(|(base, exp)| Expr::pow(base, Expr::int(exp as i64))),
        ]
    })
}

// ── Eval helpers ─────────────────────────────────────────────────

/// Build a variable map for evaluation with the given x, y, z values.
fn make_vars(x: f64, y: f64, z: f64) -> HashMap<Symbol, f64> {
    let mut m = HashMap::new();
    m.insert(Symbol::new("x"), x);
    m.insert(Symbol::new("y"), y);
    m.insert(Symbol::new("z"), z);
    m
}

/// Try to evaluate an expression; returns None if it errors (domain, overflow, etc.)
fn try_eval(expr: &Expr, vars: &HashMap<Symbol, f64>) -> Option<f64> {
    match eval(expr, vars) {
        Ok(v) if v.is_finite() => Some(v),
        _ => None,
    }
}

/// Check if two f64 values are approximately equal (relative + absolute tolerance).
fn approx_eq(a: f64, b: f64) -> bool {
    if a == b {
        return true;
    }
    let diff = (a - b).abs();
    let mag = a.abs().max(b.abs()).max(1.0);
    diff < 1e-6 * mag + 1e-10
}

// ── Property Tests ───────────────────────────────────────────────

proptest! {
    #![proptest_config(ProptestConfig::with_cases(PROPTEST_CASES))]

    // 1. Simplification preserves numeric value
    // Uses arb_polynomial to avoid Rational64 overflow on nested Pow
    #[test]
    fn simplify_preserves_value(
        expr in arb_polynomial(3),
        x in -2.0..2.0f64,
        y in -2.0..2.0f64,
        z in -2.0..2.0f64,
    ) {
        let vars = make_vars(x, y, z);
        let simplified = simplify(&expr);

        if let (Some(orig_val), Some(simp_val)) =
            (try_eval(&expr, &vars), try_eval(&simplified, &vars))
        {
            prop_assert!(
                approx_eq(orig_val, simp_val),
                "simplify changed value: orig={orig_val}, simplified={simp_val}, expr={expr}"
            );
        }
    }

    // 2. Simplification is idempotent
    // Uses arb_polynomial to avoid Rational64 overflow on nested Pow
    #[test]
    fn simplify_is_idempotent(expr in arb_polynomial(3)) {
        let once = simplify(&expr);
        let twice = simplify(&once);
        prop_assert!(
            *once == *twice,
            "simplify is not idempotent: once={}, twice={}", once, twice
        );
    }

    // 3. Parse round-trip: parse(display(e)) ≈ e (by numeric eval)
    #[test]
    fn parse_roundtrip(
        expr in arb_polynomial(2),
        x in 0.5..2.0f64,
        y in 0.5..2.0f64,
    ) {
        let displayed = format!("{}", expr);
        if let Ok(reparsed) = parse(&displayed) {
            let vars = make_vars(x, y, 1.0);
            if let (Some(orig_val), Some(re_val)) =
                (try_eval(&expr, &vars), try_eval(&reparsed, &vars))
            {
                prop_assert!(
                    approx_eq(orig_val, re_val),
                    "parse roundtrip mismatch: displayed={displayed:?}, \
                     orig={orig_val}, reparsed={re_val}"
                );
            }
        }
        // If parse fails on some exotic display format, that's OK to skip for now.
    }

    // 4. Differentiation correctness (numeric derivative check)
    #[test]
    fn differentiation_correctness(
        expr in arb_polynomial(2),
        x_val in 0.5..2.0f64,
    ) {
        let x_sym = Symbol::new("x");
        let deriv = symclaw_core::differentiate::differentiate(&expr, x_sym);
        let h = 1e-7;
        let vars_at = make_vars(x_val, 1.0, 1.0);
        let vars_plus = make_vars(x_val + h, 1.0, 1.0);
        let vars_minus = make_vars(x_val - h, 1.0, 1.0);
        if let (Some(d_val), Some(f_plus), Some(f_minus)) = (
            try_eval(&deriv, &vars_at),
            try_eval(&expr, &vars_plus),
            try_eval(&expr, &vars_minus),
        ) {
            let numeric = (f_plus - f_minus) / (2.0 * h);
            // Wider tolerance for numeric derivative (central difference has O(h²) error)
            let diff = (d_val - numeric).abs();
            let mag = d_val.abs().max(numeric.abs()).max(1.0);
            prop_assert!(
                diff < 1e-4 * mag + 1e-6,
                "derivative mismatch: symbolic={d_val}, numeric={numeric}, diff={diff}"
            );
        }
    }

    // 5. LaTeX output is non-empty for any expression
    #[test]
    fn latex_non_empty(expr in arb_expr(3)) {
        let latex = to_latex(&expr);
        prop_assert!(!latex.is_empty(), "LaTeX output was empty for expr={expr}");
    }

    // 6. contains_var returns true for expressions that use that variable
    #[test]
    fn contains_var_when_present(
        expr in arb_expr(2),
    ) {
        let x_sym = Symbol::new("x");
        let y_sym = Symbol::new("y");
        let z_sym = Symbol::new("z");

        let free = expr.free_vars();

        // If free_vars says x is free, contains_var must agree
        if free.contains(&x_sym) {
            prop_assert!(
                expr.contains_var(x_sym),
                "free_vars contains x but contains_var(x) is false: {expr}"
            );
        }
        if free.contains(&y_sym) {
            prop_assert!(
                expr.contains_var(y_sym),
                "free_vars contains y but contains_var(y) is false: {expr}"
            );
        }
        if free.contains(&z_sym) {
            prop_assert!(
                expr.contains_var(z_sym),
                "free_vars contains z but contains_var(z) is false: {expr}"
            );
        }

        // Converse: if contains_var is true, free_vars should contain it
        if expr.contains_var(x_sym) {
            prop_assert!(
                free.contains(&x_sym),
                "contains_var(x) is true but free_vars doesn't contain x: {expr}"
            );
        }
    }

    // 7. is_constant iff free_vars is empty
    #[test]
    fn is_constant_iff_no_free_vars(expr in arb_expr(3)) {
        let is_const = expr.is_constant();
        let free_empty = expr.free_vars().is_empty();
        prop_assert!(
            is_const == free_empty,
            "is_constant={} but free_vars().is_empty()={} for {}", is_const, free_empty, expr
        );
    }

    // 8. Negation is self-inverse: simplify(neg(neg(e))) ≈ simplify(e) by eval
    // Uses arb_polynomial to avoid Rational64 overflow on nested Pow
    #[test]
    fn negation_self_inverse(
        expr in arb_polynomial(3),
        x in -2.0..2.0f64,
        y in -2.0..2.0f64,
        z in -2.0..2.0f64,
    ) {
        let double_neg = Expr::neg(Expr::neg(expr.clone()));
        let simp_orig = simplify(&expr);
        let simp_dneg = simplify(&double_neg);

        let vars = make_vars(x, y, z);
        if let (Some(v1), Some(v2)) =
            (try_eval(&simp_orig, &vars), try_eval(&simp_dneg, &vars))
        {
            prop_assert!(
                approx_eq(v1, v2),
                "neg(neg(e)) != e: orig={v1}, double_neg={v2}, expr={expr}"
            );
        }
    }

    // 9. Simplification of polynomial preserves value (extra coverage)
    #[test]
    fn polynomial_simplify_preserves_value(
        expr in arb_polynomial(3),
        x in -3.0..3.0f64,
        y in -3.0..3.0f64,
    ) {
        let vars = make_vars(x, y, 1.0);
        let simplified = simplify(&expr);

        if let (Some(orig_val), Some(simp_val)) =
            (try_eval(&expr, &vars), try_eval(&simplified, &vars))
        {
            prop_assert!(
                approx_eq(orig_val, simp_val),
                "polynomial simplify changed value: orig={orig_val}, simp={simp_val}"
            );
        }
    }

    // 10. LaTeX output length is reasonable (no exponential blowup)
    #[test]
    fn latex_bounded_length(expr in arb_expr(3)) {
        let latex = to_latex(&expr);
        // A depth-3 expr should not produce megabytes of LaTeX
        prop_assert!(
            latex.len() < 100_000,
            "LaTeX output suspiciously large: {} bytes",
            latex.len()
        );
    }

    // 11. Integration correctness: d/dx(∫ f dx) ≈ f numerically
    #[test]
    fn integration_derivative_roundtrip(
        expr in arb_polynomial(2),
        x_val in 0.5..2.0f64,
    ) {
        let x_sym = Symbol::new("x");
        if let Some(antideriv) = integrate(&expr, x_sym) {
            let deriv = differentiate(&antideriv, x_sym);
            let deriv_simplified = simplify(&deriv);
            let vars = make_vars(x_val, 1.0, 1.0);
            if let (Some(orig_val), Some(roundtrip_val)) =
                (try_eval(&expr, &vars), try_eval(&deriv_simplified, &vars))
            {
                let diff = (orig_val - roundtrip_val).abs();
                let mag = orig_val.abs().max(roundtrip_val.abs()).max(1.0);
                prop_assert!(
                    diff < 1e-4 * mag + 1e-6,
                    "d/dx(∫ f dx) != f: orig={orig_val}, roundtrip={roundtrip_val}, \
                     expr={expr}"
                );
            }
        }
        // If integration returns None, skip (not all polynomials may integrate cleanly)
    }

    // 12. Solve correctness: for ax+b, solve returns x = -b/a, substitution ≈ 0
    #[test]
    fn solve_linear_equation(
        a in 1..=10i64,
        b in -10..=10i64,
    ) {
        let x_sym = Symbol::new("x");
        // Build ax + b
        let expr = Expr::add(vec![
            Expr::mul(vec![Expr::int(a), Expr::sym("x")]),
            Expr::int(b),
        ]);
        let solutions = solve(&expr, x_sym);
        if !solutions.is_empty() {
            // Verify first solution: substituting back should give ≈ 0
            let sol = &solutions[0];
            let mut vars = HashMap::new();
            if let Some(sol_val) = try_eval(sol, &HashMap::new()) {
                vars.insert(Symbol::new("x"), sol_val);
                if let Some(result_val) = try_eval(&expr, &vars) {
                    prop_assert!(
                        result_val.abs() < 1e-8,
                        "solve gave x={sol_val} but f(x)={result_val} != 0"
                    );
                }
            }
            // Also check expected value -b/a
            let expected = -(b as f64) / (a as f64);
            if let Some(sol_val) = try_eval(sol, &HashMap::new()) {
                prop_assert!(
                    approx_eq(sol_val, expected),
                    "solve: expected x={expected}, got x={sol_val}"
                );
            }
        }
    }

    // 13. Taylor approximation: for polynomial of degree d, taylor(expr, x, 0, d) ≈ expr
    #[test]
    fn taylor_polynomial_identity(
        expr in arb_polynomial(2),
        x_val in -1.0..1.0f64,
    ) {
        let x_sym = Symbol::new("x");
        let zero = Expr::int(0);
        // Simplify first to flatten powers like (x^3)^3 → x^9, then use high degree
        let expr_simplified = simplify(&expr);
        let taylor_approx = taylor(&expr_simplified, x_sym, &zero, 20);
        let vars = make_vars(x_val, 1.0, 1.0);
        // Only test single-variable polynomials (no y)
        if !expr.contains_var(Symbol::new("y"))
            && let (Some(orig_val), Some(taylor_val)) =
                (try_eval(&expr, &vars), try_eval(&taylor_approx, &vars))
            {
                let diff = (orig_val - taylor_val).abs();
                let mag = orig_val.abs().max(taylor_val.abs()).max(1.0);
                prop_assert!(
                    diff < 1e-3 * mag + 1e-6,
                    "taylor mismatch: orig={orig_val}, taylor={taylor_val}, expr={expr}"
                );
            }
    }

    // 14. Cache consistency: simplify_cached(e) == simplify(e)
    #[test]
    fn cache_consistency(expr in arb_polynomial(3)) {
        let cached = simplify_cached(&expr);
        let direct = simplify(&expr);
        prop_assert!(
            *cached == *direct,
            "cache inconsistency: cached={cached}, direct={direct}"
        );
    }

    // 15. LaTeX round-trip stability: calling twice gives identical output
    #[test]
    fn latex_stability(expr in arb_expr(3)) {
        let simplified = simplify(&expr);
        let latex1 = to_latex(&simplified);
        let latex2 = to_latex(&simplified);
        prop_assert!(
            latex1 == latex2,
            "LaTeX not stable: first={latex1:?}, second={latex2:?}"
        );
    }
}