symplex 0.17.0

//! Tests for second-order constant-coefficient nonhomogeneous ODE solver
//! using the method of undetermined coefficients.
//!
//! Covers:
//! - Constant forcing: y'' + y = 1
//! - Linear forcing: y'' + y' + y = x
//! - Quadratic forcing: y'' + y = x²
//! - Numerical verification by substitution
//! - Regression: homogeneous and first-order still work
//! - Classification of nonhomogeneous ODEs
//! - Graceful failure for non-polynomial forcing (e.g., exp)

use symplex::ode::OdeType;
use symplex::prelude::*;

// ═══════════════════════════════════════════════════════════════════════════
// Helpers
// ═══════════════════════════════════════════════════════════════════════════

/// Numerically verify a second-order ODE solution by substituting
/// C1=1, C2=1 into the solution, computing derivatives, replacing
/// formal derivatives and y in the ODE, and evaluating at a sample point.
fn verify_second_order_numerically(
    ode_expr: &Ex,
    solution: &Ex,
    constants: &[Ex],
    y: &Ex,
    x: &Ex,
    sample_x_num: i64,
    sample_x_den: i64,
) {
    let ctx = ode_expr.context();
    let one = ctx.int(1);
    let mut concrete_sol = solution.clone();
    for c in constants {
        concrete_sol = concrete_sol.subs(c, &one);
    }

    let sol_prime = concrete_sol.diff(x);
    let sol_double_prime = sol_prime.diff(x);

    let dy_formal = y.formal_diff(x);
    let d2y_formal = dy_formal.formal_diff(x);

    let residual = ode_expr
        .subs(&d2y_formal, &sol_double_prime)
        .subs(&dy_formal, &sol_prime)
        .subs(y, &concrete_sol);

    let sample_val = ctx.rational(sample_x_num, sample_x_den);
    let residual_at = residual.subs(x, &sample_val);

    let val = residual_at
        .eval_f64()
        .expect("evalf_f64 should succeed for second-order ODE residual");
    assert!(
        val.abs() < 1e-4,
        "Second-order ODE residual should be ~0, got {val} at x={sample_x_num}/{sample_x_den}\n  \
         solution (C=1): {concrete_sol}\n  residual: {residual_at}"
    );
}

/// Numerically verify a first-order ODE solution.
fn verify_first_order_numerically(
    ode_expr: &Ex,
    solution: &Ex,
    constants: &[Ex],
    y: &Ex,
    x: &Ex,
    sample_x_num: i64,
    sample_x_den: i64,
) {
    let ctx = ode_expr.context();
    let one = ctx.int(1);
    let mut concrete_sol = solution.clone();
    for c in constants {
        concrete_sol = concrete_sol.subs(c, &one);
    }

    let sol_prime = concrete_sol.diff(x);

    let dy_formal = y.formal_diff(x);
    let residual = ode_expr.subs(&dy_formal, &sol_prime).subs(y, &concrete_sol);

    let sample_val = ctx.rational(sample_x_num, sample_x_den);
    let residual_at = residual.subs(x, &sample_val);

    let val = residual_at
        .eval_f64()
        .expect("evalf_f64 should succeed for first-order ODE residual");
    assert!(
        val.abs() < 1e-6,
        "First-order ODE residual should be ~0, got {val} at x={sample_x_num}/{sample_x_den}\n  \
         solution (C=1): {concrete_sol}\n  residual: {residual_at}"
    );
}

// ═══════════════════════════════════════════════════════════════════════════
// Test 1: Constant forcing  y'' + y = 1
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_second_order_constant_rhs() {
    let ctx = Context::new();
    // y'' + y - 1 = 0  ⟹  y'' + y = 1
    // Particular solution: y_p = 1 (since c=1 ≠ 0, y_p = k/c = 1/1)
    // Homogeneous: y_h = C1·exp(ix) + C2·exp(-ix)
    // General: y = C1·exp(ix) + C2·exp(-ix) + 1
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let one = ctx.int(1);
    let ode = &d2y + &y - &one; // y'' + y - 1 = 0

    let sol = ode.solve_ode(&y, &x);
    assert!(!sol.has_unevaluated(), "y'' + y = 1 should be solvable");

    let s = format!("{sol}");
    assert!(
        s.contains("C1") && s.contains("C2"),
        "should have two constants: {s}"
    );

    // The solution should contain exp terms (homogeneous) and the constant 1 (particular).
    // Verify numerically at a test point.
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 7, 10);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Test 2: Linear forcing  y'' + y' + y = x
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_second_order_linear_rhs() {
    let ctx = Context::new();
    // y'' + y' + y - x = 0  ⟹  y'' + y' + y = x
    // With b=1, c=1, rhs = x  (degree 1)
    // Try y_p = Ax + B:  y_p''=0, y_p'=A
    //   A + Ax + B = x  ⟹  cA = 1 → A=1, bA + cB = 0 → 1 + B = 0 → B=-1
    //   y_p = x - 1
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let ode = &d2y + &dy + &y - &x; // y'' + y' + y - x = 0

    let sol = ode.solve_ode(&y, &x);
    assert!(
        !sol.has_unevaluated(),
        "y'' + y' + y = x should be solvable"
    );

    let s = format!("{sol}");
    assert!(
        s.contains("C1") && s.contains("C2"),
        "should have two constants: {s}"
    );

    // Verify numerically
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 1, 2);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Test 3: Quadratic forcing  y'' + y = x²
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_second_order_quadratic_rhs() {
    let ctx = Context::new();
    // y'' + y - x² = 0  ⟹  y'' + y = x²
    // With b=0, c=1, rhs = x² (degree 2)
    // Try y_p = Ax² + Bx + D:
    //   y_p'' = 2A, y_p' = 2Ax + B
    //   2A + Ax² + Bx + D = x²
    //   A = 1, B = 0, 2A + D = 0 → D = -2
    //   y_p = x² - 2
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let x_sq = x.powi(2);
    let ode = &d2y + &y - &x_sq; // y'' + y - x² = 0

    let sol = ode.solve_ode(&y, &x);
    assert!(!sol.has_unevaluated(), "y'' + y = x² should be solvable");

    let s = format!("{sol}");
    assert!(
        s.contains("C1") && s.contains("C2"),
        "should have two constants: {s}"
    );

    // Verify numerically
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 3, 10);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Test 4: Numerical verification by substitution
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_second_order_nonhomogeneous_verify() {
    let ctx = Context::new();
    // y'' - 3y' + 2y = 6  (distinct real roots r=1,2)
    // Particular: y_p = 6/2 = 3
    // Homogeneous: y_h = C1*exp(x) + C2*exp(2x)
    // General: y = C1*exp(x) + C2*exp(2x) + 3
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let six = ctx.int(6);
    let ode = &d2y - &(&dy * 3) + &(&y * 2) - &six; // y'' - 3y' + 2y - 6 = 0

    let sol = ode
        .try_solve_ode(&y, &x)
        .expect("y'' - 3y' + 2y = 6 should be solvable");

    // Verify at multiple points for robustness
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 1, 4);
    }
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 1, 2);
    }
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 3, 4);
    }

    // Also verify with C1=0, C2=0 to isolate the particular solution
    let c1 = ctx.symbol("C1");
    let c2 = ctx.symbol("C2");
    let zero = ctx.int(0);
    let particular = sol.subs(&c1, &zero).subs(&c2, &zero);
    let particular_s = format!("{particular}");
    // The particular solution should simplify to 3
    let particular_val = particular.eval_f64();
    if let Ok(v) = particular_val {
        assert!(
            (v - 3.0).abs() < 1e-10,
            "particular solution should be 3, got {v} ({particular_s})"
        );
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Test 5: Homogeneous still works (regression)
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_homogeneous_still_works() {
    let ctx = Context::new();
    // y'' + y = 0 should still be solved by the homogeneous path
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let ode = &d2y + &y; // y'' + y = 0

    let sol = ode.solve_ode(&y, &x);
    // This should still solve (complex roots ±i)
    if !sol.has_unevaluated() {
        let s = format!("{sol}");
        assert!(
            s.contains("C1") && s.contains("C2"),
            "homogeneous should have two constants: {s}"
        );
    }

    // Another homogeneous: y'' - 3y' + 2y = 0
    let ode2 = &d2y - &(&dy * 3) + &(&y * 2);
    let sol2 = ode2
        .try_solve_ode(&y, &x)
        .expect("y'' - 3y' + 2y = 0 should still be solvable");
    let s2 = format!("{sol2}");
    assert!(
        s2.contains("C1") && s2.contains("C2"),
        "homogeneous should have two constants: {s2}"
    );
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode2, &sol2, &[c1, c2], &y, &x, 3, 10);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Test 6: classify_ode returns correct type for nonhomogeneous
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_classify_nonhomogeneous() {
    let ctx = Context::new();
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);

    // y'' + y - 1 = 0 should be classified as nonhomogeneous
    let one = ctx.int(1);
    let ode = &d2y + &y - &one;
    assert_eq!(
        ode.classify_ode(&y, &x),
        OdeType::SecondOrderLinearCCNonHomogeneous,
        "y'' + y = 1 should classify as nonhomogeneous"
    );

    // y'' - 3y' + 2y - x = 0 should be classified as nonhomogeneous
    let ode2 = &d2y - &(&dy * 3) + &(&y * 2) - &x;
    assert_eq!(
        ode2.classify_ode(&y, &x),
        OdeType::SecondOrderLinearCCNonHomogeneous,
        "y'' - 3y' + 2y = x should classify as nonhomogeneous"
    );

    // y'' + y = 0 should still classify as homogeneous
    let ode3 = &d2y + &y;
    assert_eq!(
        ode3.classify_ode(&y, &x),
        OdeType::SecondOrderLinearCCHomogeneous,
        "y'' + y = 0 should classify as homogeneous"
    );
}

// ═══════════════════════════════════════════════════════════════════════════
// Test 7: Non-polynomial forcing returns None gracefully
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_exponential_rhs() {
    let ctx = Context::new();
    // y'' + y = exp(x) — exponential forcing is NOT a polynomial,
    // so undetermined coefficients for polynomial rhs should not apply.
    // The solver should return None (or possibly handle it in the future).
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let ode = &d2y + &y - &x.exp(); // y'' + y - exp(x) = 0

    let sol = ode.solve_ode(&y, &x);
    // It's OK if this returns unevaluated — we just must not return a wrong answer.
    // If it does return something, verify it numerically to ensure correctness.
    if !sol.has_unevaluated() {
        {
            let c1 = ctx.symbol("C1");
            let c2 = ctx.symbol("C2");
            verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 1, 2);
        }
    }
    // No assertion failure = pass (graceful None or correct answer)
}

// ═══════════════════════════════════════════════════════════════════════════
// Test 8: First-order still works (regression)
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_first_order_still_works() {
    let ctx = Context::new();
    // Regression: y' + y = 0 still solves correctly
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let ode = &dy + &y; // y' + y = 0

    let sol = ode
        .try_solve_ode(&y, &x)
        .expect("y' + y = 0 should still be solvable");
    let s = format!("{sol}");
    assert!(s.contains("C1"), "should have constant: {s}");
    assert!(s.contains("exp"), "should contain exp: {s}");

    {
        let c1 = ctx.symbol("C1");
        verify_first_order_numerically(&ode, &sol, &[c1], &y, &x, 1, 2);
    }

    // Also: y' = x should still work (simple separable)
    let ode2 = &dy - &x; // y' - x = 0
    let sol2 = ode2
        .try_solve_ode(&y, &x)
        .expect("y' = x should still be solvable");
    let s2 = format!("{sol2}");
    assert!(s2.contains("C1"), "should have constant: {s2}");

    {
        let c1 = ctx.symbol("C1");
        verify_first_order_numerically(&ode2, &sol2, &[c1], &y, &x, 3, 2);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Additional: distinct real roots with linear forcing
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_distinct_roots_linear_forcing() {
    let ctx = Context::new();
    // y'' - 3y' + 2y = x  (roots r=1, r=2)
    // Particular: y_p = Ax + B
    //   0 - 3A + 2(Ax + B) = x  ⟹  2A = 1 → A=1/2, -3A + 2B = 0 → B=3/4
    //   y_p = x/2 + 3/4
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let ode = &d2y - &(&dy * 3) + &(&y * 2) - &x; // y'' - 3y' + 2y - x = 0

    let sol_expr = ode
        .try_solve_ode(&y, &x)
        .expect("y'' - 3y' + 2y = 6 should be solvable");
    let s = format!("{sol_expr}");
    assert!(
        s.contains("C1") && s.contains("C2"),
        "should have two constants: {s}"
    );

    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol_expr, &[c1, c2], &y, &x, 1, 4);
    }
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol_expr, &[c1, c2], &y, &x, 1, 1);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Additional: repeated root with constant forcing
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_repeated_root_constant_forcing() {
    let ctx = Context::new();
    // y'' - 2y' + y = 4  (repeated root r=1)
    // Particular: y_p = 4/1 = 4  (since c=1)
    // Homogeneous: y_h = (C1 + C2*x)*exp(x)
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let four = ctx.int(4);
    let ode = &d2y - &(&dy * 2) + &y - &four; // y'' - 2y' + y - 4 = 0

    let sol_expr = ode
        .try_solve_ode(&y, &x)
        .expect("y'' - 2y' + y = exp(x) should be solvable");
    let s = format!("{sol_expr}");
    assert!(
        s.contains("C1") && s.contains("C2"),
        "should have two constants: {s}"
    );

    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol_expr, &[c1, c2], &y, &x, 1, 5);
    }
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol_expr, &[c1, c2], &y, &x, 3, 10);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Additional: c = 0 case (y'' + by' = f(x))
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_c_zero_constant_forcing() {
    let ctx = Context::new();
    // y'' + y' - 2 = 0  ⟹  y'' + y' = 2
    // c=0, b=1, rhs = constant 2
    // Case 2: y_p = B_0*x → y_p' = B_0, y_p'' = 0 → b*B_0 = 2 → B_0 = 2
    // y_p = 2x
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let two = ctx.int(2);
    let ode = &d2y + &dy - &two; // y'' + y' - 2 = 0

    let sol = ode.solve_ode(&y, &x);
    assert!(!sol.has_unevaluated(), "y'' + y' = 2 should be solvable");
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 1, 2);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Additional: c = 0, b = 0 case (y'' = f(x))
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_b_c_zero_constant_forcing() {
    let ctx = Context::new();
    // y'' - 6 = 0  ⟹  y'' = 6
    // c=0, b=0, rhs = 6
    // Case 3: y_p = 6/(1*2) * x² = 3x²
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let six = ctx.int(6);
    let ode = &d2y - &six; // y'' - 6 = 0

    let sol = ode.solve_ode(&y, &x);
    assert!(!sol.has_unevaluated(), "y'' = 6 should be solvable");
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 1, 2);
    }
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 2, 1);
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Additional: checkodesol integration
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_nonhomogeneous_checkodesol() {
    let ctx = Context::new();
    // y'' - 3y' + 2y = 6 → particular y_p = 3
    // Check that a known particular solution passes checkodesol
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let six = ctx.int(6);
    let ode = &d2y - &(&dy * 3) + &(&y * 2) - &six;

    // y = 3 should be a particular solution
    let particular = ctx.int(3);
    let ok = ode.check_ode_solution(&particular, &y, &x);
    assert!(ok, "y=3 should satisfy y'' - 3y' + 2y = 6");

    // y = 0 should NOT be a solution
    let wrong = ctx.int(0);
    let not_ok = ode.check_ode_solution(&wrong, &y, &x);
    assert!(!not_ok, "y=0 should NOT satisfy y'' - 3y' + 2y = 6");
}

// ═══════════════════════════════════════════════════════════════════════════
// Additional: scaled leading coefficient
// ═══════════════════════════════════════════════════════════════════════════

#[test]
fn ode_scaled_leading_coefficient() {
    let ctx = Context::new();
    // 2y'' + 2y = 4  ⟹  y'' + y = 2  ⟹  y_p = 2
    let x = ctx.symbol("x");
    let y = ctx.symbol("y");
    let dy = y.formal_diff(&x);
    let d2y = dy.formal_diff(&x);
    let two = ctx.int(2);
    let four = ctx.int(4);
    let ode = &(&d2y * 2) + &(&y * 2) - &four; // 2y'' + 2y - 4 = 0

    let sol = ode.solve_ode(&y, &x);
    assert!(
        !sol.has_unevaluated(),
        "2y'' + 2y = 4 should be solvable (normalises to y'' + y = 2)"
    );

    let _ = two; // suppress unused warning
    {
        let c1 = ctx.symbol("C1");
        let c2 = ctx.symbol("C2");
        verify_second_order_numerically(&ode, &sol, &[c1, c2], &y, &x, 1, 3);
    }
}