# ============================================================
# File: _solve_numerical_ode_methods_advanced.py
# Topic: Numerical ODE Methods Beyond RK4
#
# Covers "textbook + app-friendly" templates for:
#   1) Fixed-step one-step methods:
#        - Forward Euler
#        - Heun / Improved Euler (explicit trapezoid)
#        - Midpoint (RK2)
#        - RK4 (reference)
#   2) Adaptive RK (RK45 / Dormand–Prince idea):
#        - error estimate and step-size control formula
#   3) Multistep methods:
#        - Adams–Bashforth (AB2)
#        - Adams–Moulton (AM2 / trapezoid) + predictor–corrector
#   4) Stiff solvers (implicit):
#        - Backward Euler
#        - Implicit trapezoid
#        - BDF2 (template)
#
# We generate:
#   - the update formulas
#   - stability / error-control notes
#   - pseudocode-style steps (no heavy numeric integration here)
#
# Output:
#   { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
#
# Notes:
#   - Avoid f-strings to prevent brace issues with LaTeX.
#   - SymPy is used for parsing/latex only.
# ============================================================

from __future__ import annotations

from typing import Optional, Union, Dict, Any, Literal

import sympy as sp


NumericalODEKind = Literal[
    "euler_forward",
    "heun_improved_euler",
    "midpoint_rk2",
    "rk4_reference",
    "rk45_adaptive_template",
    "adams_bashforth_ab2",
    "predictor_corrector_ab2_am2",
    "backward_euler_implicit",
    "implicit_trapezoid",
    "bdf2_template",
    "stiffness_stability_test_lambda",
]


def solve_numerical_ode_methods_advanced(
    *,
    kind: NumericalODEKind,

    # Symbols / placeholders
    t_symbol: str = "t",
    y_symbol: str = "y",
    h_symbol: str = "h",
    f_symbol: str = "f",

    # Optional: stiffness test lambda for y' = λ y
    lam_: Union[int, float, str, sp.Expr] = -10,

    # Adaptive control parameters (template only)
    tol_: Union[int, float, str, sp.Expr] = sp.Rational(1, 6),
    p_control_: int = 4,  # typical for RK45 error control uses order 4 embedded estimate

    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    t = sp.Symbol(t_symbol, real=True)
    h = sp.Symbol(h_symbol, positive=True, real=True)
    lam = sp.sympify(lam_, locals={t_symbol: t, h_symbol: h})

    tol = sp.sympify(tol_, locals={t_symbol: t, h_symbol: h})

    # We keep f as a symbolic function f(t,y)
    y = sp.Function(y_symbol)
    f = sp.Function(f_symbol)

    base_locals = {
        t_symbol: t,
        h_symbol: h,
        y_symbol: y,
        f_symbol: f,
        "pi": sp.pi,
        "E": sp.E,
        "exp": sp.exp,
        "log": sp.log,
        "ln": sp.log,
        "sqrt": sp.sqrt,
        "sin": sp.sin,
        "cos": sp.cos,
        "Abs": sp.Abs,
        "Piecewise": sp.Piecewise,
        "Heaviside": sp.Heaviside,
    }
    if locals_map:
        base_locals.update(locals_map)

    def Ltx(expr: sp.Expr) -> str:
        return sp.latex(expr)

    def card(title: str, body_aligned: str, note: str = "") -> str:
        note_html = ("<div class='note'>" + note + "</div>") if note else ""
        return (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>" + title + "<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            + body_aligned +
            "\\end{aligned}\\)"
            + note_html +
            "</div></div>"
        )

    if simplify:
        lam = sp.simplify(lam)
        tol = sp.simplify(tol)

    # Generic nodes for notation
    tn = sp.Symbol("t_n", real=True)
    yn = sp.Symbol("y_n", real=True)
    tn1 = sp.Symbol("t_{n+1}", real=True)
    yn1 = sp.Symbol("y_{n+1}", real=True)
    fn = sp.Symbol("f_n", real=True)
    fnm1 = sp.Symbol("f_{n-1}", real=True)

    # =========================================================
    # 1) Forward Euler
    # =========================================================
    if kind == "euler_forward":
        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle t_{n+1}=t_n+h \\\\[6pt]"
            "& \\displaystyle \\Rightarrow y_{n+1}=y_n+h\\,f(t_n,y_n) \\\\[10pt]"
            "& \\displaystyle \\text{Order: }O(h)\\;\\text{(global)} \\\\[4pt]"
            "& \\displaystyle \\text{Stability note (stiffness): may require very small }h. \\\\"
        )
        soln_s = card(
            "Numerical ODE — Forward Euler (1st Order)",
            body,
            "Forward Euler is simple but can be unstable for stiff problems unless h is very small.",
        )
        problem_s = "Write the Forward Euler method for \\(y'=f(t,y)\\)."
        answer_s = "\\(y_{n+1}=y_n+h f(t_n,y_n)\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 2) Heun / Improved Euler (explicit trapezoid)
    # =========================================================
    if kind == "heun_improved_euler":
        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle \\text{Predict: } \\tilde y_{n+1}=y_n+h f(t_n,y_n) \\\\[6pt]"
            "& \\displaystyle \\text{Correct: } y_{n+1}=y_n+\\frac{h}{2}\\Big(f(t_n,y_n)+f(t_{n+1},\\tilde y_{n+1})\\Big) \\\\[10pt]"
            "& \\displaystyle \\text{Order: }O(h^2)\\;\\text{(global)} \\\\"
        )
        soln_s = card(
            "Numerical ODE — Heun / Improved Euler (2nd Order)",
            body,
            "Heun is a predictor–corrector RK2 method; much better accuracy than Euler for similar cost.",
        )
        problem_s = "Write Heun’s (Improved Euler) method for \\(y'=f(t,y)\\)."
        answer_s = "\\(\\tilde y_{n+1}=y_n+h f(t_n,y_n),\\; y_{n+1}=y_n+\\tfrac{h}{2}(f_n+f(t_{n+1},\\tilde y_{n+1}))\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 3) Midpoint (RK2)
    # =========================================================
    if kind == "midpoint_rk2":
        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle k_1=f(t_n,y_n) \\\\[4pt]"
            "& \\displaystyle k_2=f\\left(t_n+\\frac{h}{2},\\; y_n+\\frac{h}{2}k_1\\right) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow y_{n+1}=y_n+h\\,k_2 \\\\[10pt]"
            "& \\displaystyle \\text{Order: }O(h^2)\\;\\text{(global)} \\\\"
        )
        soln_s = card(
            "Numerical ODE — Midpoint Method (RK2)",
            body,
            "Midpoint RK2 samples the slope at the middle of the step; usually more accurate than Heun for some problems.",
        )
        problem_s = "Write the Midpoint (RK2) method for \\(y'=f(t,y)\\)."
        answer_s = "\\(k_1=f(t_n,y_n),\\; k_2=f(t_n+\\tfrac h2, y_n+\\tfrac h2 k_1),\\; y_{n+1}=y_n+h k_2\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 4) RK4 reference (for completeness)
    # =========================================================
    if kind == "rk4_reference":
        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle k_1=f(t_n,y_n) \\\\"
            "& \\displaystyle k_2=f\\left(t_n+\\frac{h}{2},\\; y_n+\\frac{h}{2}k_1\\right) \\\\"
            "& \\displaystyle k_3=f\\left(t_n+\\frac{h}{2},\\; y_n+\\frac{h}{2}k_2\\right) \\\\"
            "& \\displaystyle k_4=f\\left(t_n+h,\\; y_n+h k_3\\right) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow y_{n+1}=y_n+\\frac{h}{6}(k_1+2k_2+2k_3+k_4) \\\\[10pt]"
            "& \\displaystyle \\text{Order: }O(h^4)\\;\\text{(global)} \\\\"
        )
        soln_s = card(
            "Numerical ODE — RK4 (Reference)",
            body,
            "RK4 is accurate and robust for smooth non-stiff problems, but it is fixed-step unless paired with step-size control.",
        )
        problem_s = "Write the RK4 update formula for \\(y'=f(t,y)\\)."
        answer_s = "\\(y_{n+1}=y_n+\\tfrac{h}{6}(k_1+2k_2+2k_3+k_4)\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 5) Adaptive RK45 template (error estimate + step control)
    # =========================================================
    if kind == "rk45_adaptive_template":
        p = int(p_control_)
        # Step-size update formula (template):
        # h_new = h * (tol / e)^(1/(p+1))
        e = sp.Symbol("e_n", positive=True, real=True)
        h_new = h * (tol / e) ** (sp.Rational(1, p + 1))
        if simplify:
            h_new = sp.simplify(h_new)

        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle \\text{Embedded pair computes } y_{n+1}^{(5)} \\text{ and } y_{n+1}^{(4)} \\text{ in one step.} \\\\[6pt]"
            "& \\displaystyle e_n = \\left|y_{n+1}^{(5)}-y_{n+1}^{(4)}\\right| \\quad (\\text{error estimate}) \\\\[10pt]"
            "& \\displaystyle \\text{Accept if } e_n\\le \\text{tol}.\\;\\; \\text{Otherwise reject and reduce }h. \\\\[6pt]"
            "& \\displaystyle \\text{Step control (typical): } h_{new}=h\\left(\\frac{\\text{tol}}{e_n}\\right)^{\\frac{1}{p+1}},\\; p=" + str(p) + " \\\\[6pt]"
            "& \\displaystyle \\Rightarrow h_{new}=" + Ltx(h_new) + " \\\\[10pt]"
            "& \\displaystyle \\text{Pseudocode:} \\\\"
            "& \\displaystyle \\text{while }t<T:\\;\\text{compute }(y^{(5)},y^{(4)}),\\; e=|\\cdot| \\\\"
            "& \\displaystyle \\quad \\text{if }e\\le tol:\\; t\\leftarrow t+h,\\; y\\leftarrow y^{(5)} \\\\"
            "& \\displaystyle \\quad h\\leftarrow h\\cdot(tol/e)^{1/(p+1)} \\;\\text{(with safety factors/clamps)} \\\\"
        )
        soln_s = card(
            "Numerical ODE — Adaptive RK45 (Template: Error + Step-Size Control)",
            body,
            "This is the standard ‘automatic step size’ engine used by many practical solvers (Dormand–Prince).",
        )
        problem_s = "Explain adaptive RK45 via an embedded error estimate and step-size control."
        answer_s = "\\(e_n=|y_{n+1}^{(5)}-y_{n+1}^{(4)}|,\\; h_{new}=h(\\tfrac{tol}{e_n})^{1/(p+1)}\\) (typical)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 6) Adams–Bashforth AB2
    # =========================================================
    if kind == "adams_bashforth_ab2":
        body = (
            "& \\displaystyle y' = f(t,y),\\quad f_n=f(t_n,y_n) \\\\[6pt]"
            "& \\displaystyle \\text{AB2 (explicit): } y_{n+1}=y_n+\\frac{h}{2}(3f_n-f_{n-1}) \\\\[10pt]"
            "& \\displaystyle \\text{Order: }O(h^2)\\;\\text{(global)} \\\\[6pt]"
            "& \\displaystyle \\text{Needs startup (e.g., one RK4 step to get }y_1\\text{).} \\\\"
        )
        soln_s = card(
            "Numerical ODE — Adams–Bashforth AB2 (Explicit Multistep)",
            body,
            "Multistep methods reuse previous f-evals, but require starting values and can have stability limits.",
        )
        problem_s = "Write the AB2 (Adams–Bashforth) method."
        answer_s = "\\(y_{n+1}=y_n+\\tfrac{h}{2}(3f_n-f_{n-1})\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 7) Predictor–Corrector: AB2 predictor + AM2 corrector
    # =========================================================
    if kind == "predictor_corrector_ab2_am2":
        body = (
            "& \\displaystyle y' = f(t,y),\\quad f_n=f(t_n,y_n) \\\\[6pt]"
            "& \\displaystyle \\text{Predict (AB2): } \\tilde y_{n+1}=y_n+\\frac{h}{2}(3f_n-f_{n-1}) \\\\[6pt]"
            "& \\displaystyle \\text{Correct (AM2 / trapezoid): } y_{n+1}=y_n+\\frac{h}{2}\\Big(f_n+f(t_{n+1},\\tilde y_{n+1})\\Big) \\\\[10pt]"
            "& \\displaystyle \\text{(Optionally iterate the corrector by recomputing }f(t_{n+1},y_{n+1}).) \\\\"
        )
        soln_s = card(
            "Numerical ODE — Predictor–Corrector (AB2 + AM2)",
            body,
            "AB predicts quickly; AM corrects with better stability/accuracy. Startup values still needed.",
        )
        problem_s = "Write an AB2–AM2 predictor–corrector pair."
        answer_s = "\\(\\tilde y_{n+1}=y_n+\\tfrac{h}{2}(3f_n-f_{n-1}),\\; y_{n+1}=y_n+\\tfrac{h}{2}(f_n+f(t_{n+1},\\tilde y_{n+1}))\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 8) Backward Euler (implicit)
    # =========================================================
    if kind == "backward_euler_implicit":
        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle \\text{Backward Euler (implicit): } y_{n+1}=y_n+h\\,f(t_{n+1},y_{n+1}) \\\\[10pt]"
            "& \\displaystyle \\text{Each step requires solving } \\; y_{n+1}-h f(t_{n+1},y_{n+1})=y_n \\\\[6pt]"
            "& \\displaystyle \\text{Typically solved by Newton: } y^{(m+1)}=y^{(m)}-\\frac{F(y^{(m)})}{F'(y^{(m)})} \\\\"
        )
        soln_s = card(
            "Numerical ODE — Backward Euler (Implicit, Stiff-Friendly)",
            body,
            "Backward Euler is A-stable for linear test problems and is a baseline stiff integrator.",
        )
        problem_s = "Write the Backward Euler method and the implicit equation to solve each step."
        answer_s = "\\(y_{n+1}=y_n+h f(t_{n+1},y_{n+1})\\) (solve implicitly for \\(y_{n+1}\\))."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 9) Implicit trapezoid (A-stable)
    # =========================================================
    if kind == "implicit_trapezoid":
        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle \\text{Implicit trapezoid: } y_{n+1}=y_n+\\frac{h}{2}\\Big(f(t_n,y_n)+f(t_{n+1},y_{n+1})\\Big) \\\\[10pt]"
            "& \\displaystyle \\text{Implicit solve needed (Newton or fixed-point).} \\\\"
        )
        soln_s = card(
            "Numerical ODE — Implicit Trapezoid (A-stable)",
            body,
            "Often good for moderately stiff problems; can show oscillations for very stiff systems (no L-stability).",
        )
        problem_s = "Write the implicit trapezoidal method."
        answer_s = "\\(y_{n+1}=y_n+\\tfrac{h}{2}(f(t_n,y_n)+f(t_{n+1},y_{n+1}))\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 10) BDF2 template
    # =========================================================
    if kind == "bdf2_template":
        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle \\text{BDF2 (implicit): } \\frac{3y_{n+1}-4y_n+y_{n-1}}{2h}=f(t_{n+1},y_{n+1}) \\\\[10pt]"
            "& \\displaystyle \\Rightarrow 3y_{n+1}-2h f(t_{n+1},y_{n+1}) = 4y_n - y_{n-1} \\\\[6pt]"
            "& \\displaystyle \\text{Solve implicitly for }y_{n+1}\\text{ each step. Startup requires }y_1\\text{ (e.g., RK4).} \\\\"
        )
        soln_s = card(
            "Numerical ODE — BDF2 (Template, Stiff Solver Family)",
            body,
            "BDF methods are widely used for stiffness; they require implicit solves (often Newton + linear solves for systems).",
        )
        problem_s = "Write the BDF2 method and show the implicit equation for \\(y_{n+1}\\)."
        answer_s = "\\(\\frac{3y_{n+1}-4y_n+y_{n-1}}{2h}=f(t_{n+1},y_{n+1})\\) (implicit)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 11) Stiffness / stability test on y' = λ y
    # =========================================================
    if kind == "stiffness_stability_test_lambda":
        # Euler stability: |1+hλ|<1 -> 0<h<2/|λ| for λ<0
        cond = sp.Abs(1 + h * lam) < 1
        h_max = 2 / sp.Abs(lam) if lam != 0 else sp.oo
        if simplify:
            h_max = sp.simplify(h_max)

        # Backward Euler amplification factor: 1/(1-hλ)
        amp_be = 1 / (1 - h * lam)
        if simplify:
            amp_be = sp.simplify(amp_be)

        body = (
            "& \\displaystyle y' = \\lambda y,\\quad \\lambda=" + Ltx(lam) + " \\\\[8pt]"
            "& \\displaystyle \\text{Forward Euler: } y_{n+1}=(1+h\\lambda)y_n \\\\[4pt]"
            "& \\displaystyle \\text{Stability requires } |1+h\\lambda|<1 \\\\[4pt]"
            "& \\displaystyle \\Rightarrow h<\\frac{2}{|\\lambda|} = " + Ltx(h_max) + " \\quad (\\lambda<0) \\\\[10pt]"
            "& \\displaystyle \\text{Backward Euler: } y_{n+1}=\\frac{y_n}{1-h\\lambda} \\\\[4pt]"
            "& \\displaystyle \\text{Amplification factor: } \\frac{1}{1-h\\lambda}=" + Ltx(amp_be) + " \\\\[4pt]"
            "& \\displaystyle \\Rightarrow \\text{For }\\lambda<0,\\; \\text{stable for all }h>0\\; (A\\text{-stable}). \\\\"
        )
        soln_s = card(
            "Numerical ODE — Stiffness Stability Test (y' = λy)",
            body,
            "This explains why explicit methods need tiny h for large negative λ (stiffness), while implicit methods remain stable.",
        )
        problem_s = "Compare Euler vs backward Euler stability on the test equation \\(y'=\\lambda y\\)."
        answer_s = "Euler stable only if \\(|1+h\\lambda|<1\\); backward Euler stable for all \\(h>0\\) when \\(\\lambda<0\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError(
        "Unknown kind={!r}. Use "
        "'euler_forward','heun_improved_euler','midpoint_rk2','rk4_reference',"
        "'rk45_adaptive_template','adams_bashforth_ab2','predictor_corrector_ab2_am2',"
        "'backward_euler_implicit','implicit_trapezoid','bdf2_template','stiffness_stability_test_lambda'.".format(kind)
    )
