# ============================================================
# File: _solve_numerical_stability_and_stiffness_analysis.py
# Topic: Stability & stiffness analysis tools (meta-layer)
#
# Covers:
#   - Linear test equation y' = λ y, z = hλ
#   - Amplification factor / stability function R(z)
#   - Absolute stability condition |R(z)| <= 1
#   - Stability functions:
#       Forward Euler:     R(z)=1+z
#       Backward Euler:    R(z)=1/(1-z)    (A-stable)
#       Trapezoidal (CN):  R(z)=(1+z/2)/(1-z/2)  (A-stable, not L-stable)
#       RK4: polynomial R(z) (bounded stability region)
#   - Step-size rules for real λ<0 (Euler)
#   - Stiffness concept + stiffness ratio for y' = Ay (via eigenvalues)
#   - Practical stiffness heuristics (RK45 step collapse, rejected steps, Jacobian norm)
#   - PDE CFL rules (heat FTCS, wave centered)
#
# Output:
#   { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
#
# Notes:
#   - Avoid f-strings to prevent LaTeX brace errors in HTML strings.
# ============================================================

from __future__ import annotations

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

import sympy as sp


StabilityKind = Literal[
    "linear_test_equation_overview",
    "absolute_stability_definition",
    "forward_euler_stability_region",
    "backward_euler_a_stability",
    "trapezoidal_a_stability_not_l_stable",
    "rk4_stability_function",
    "stiffness_definition_and_ratio",
    "stiffness_detection_heuristics",
    "pde_cfl_rules",
]


def solve_numerical_stability_and_stiffness_analysis(
    *,
    kind: StabilityKind,
    # optional parameters for display
    lam_: Union[int, float, str, sp.Expr] = "-10",
    h_: Union[int, float, str, sp.Expr] = "0.05",
    kappa_: Union[int, float, str, sp.Expr] = "1000",
    # symbols
    t_symbol: str = "t",
    y_symbol: str = "y",
    h_symbol: str = "h",
    lambda_symbol: str = "\\lambda",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    t = sp.Symbol(t_symbol, real=True)
    y = sp.Function(y_symbol)
    h = sp.Symbol(h_symbol, positive=True, real=True)
    lam = sp.Symbol(lambda_symbol, complex=True)
    z = sp.Symbol("z", complex=True)

    base_locals = {
        t_symbol: t,
        y_symbol: y,
        h_symbol: h,
        "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,
        "I": sp.I,
    }
    if locals_map:
        base_locals.update(locals_map)

    def P(v: Union[int, float, str, sp.Expr]) -> sp.Expr:
        return sp.sympify(v, locals=base_locals)

    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>"
        )

    lam_val = P(lam_)
    h_val = P(h_)
    kappa_val = P(kappa_)
    if simplify:
        lam_val = sp.simplify(lam_val)
        h_val = sp.simplify(h_val)
        kappa_val = sp.simplify(kappa_val)

    # ------------------------------------------------------------
    # 1) Linear test equation overview
    # ------------------------------------------------------------
    if kind == "linear_test_equation_overview":
        # FIX: Replaced < with \lt to prevent HTML parsing errors
        body = (
            "& \\displaystyle \\text{Linear test ODE: } y' = \\lambda y,\\quad \\lambda\\in\\mathbb{C} \\\\[6pt]"
            "& \\displaystyle \\text{Exact solution: } y(t)=e^{\\lambda t}y(0) \\\\[6pt]"
            "& \\displaystyle \\text{If }\\Re(\\lambda) \\lt 0,\\; y(t)\\to 0\\text{ as }t\\to\\infty \\;\\Rightarrow\\; \\text{numerical method should decay too.} \\\\"
        )
        soln_s = card(
            "Stability Meta-Tool — The Linear Test Equation",
            body,
            "Nearly all absolute-stability theory is defined by applying a method to y' = λy and checking if the numerical solution decays when the true one does.",
        )
        problem_s = "State the linear test equation and its exact solution."
        answer_s = "\\(y'=\\lambda y\\Rightarrow y(t)=e^{\\lambda t}y(0)\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # 2) Absolute stability definition via R(z)
    # ------------------------------------------------------------
    if kind == "absolute_stability_definition":
        body = (
            "& \\displaystyle \\text{Apply a one-step method to } y' = \\lambda y. \\\\[4pt]"
            "& \\displaystyle \\text{Typically you get: } y_{n+1}=R(z)\\,y_n,\\quad z=h\\lambda \\\\[8pt]"
            "& \\displaystyle \\text{Absolute stability condition: } |R(z)|\\le 1 \\\\[6pt]"
            "& \\displaystyle \\text{Stability region: } \\{z\\in\\mathbb{C}: |R(z)|\\le 1\\} \\\\"
        )
        soln_s = card(
            "Absolute Stability — Amplification Factor",
            body,
            "R(z) is the method’s stability function. Once you know R(z), you can state stable step-size constraints and compare explicit vs implicit methods.",
        )
        problem_s = "Define absolute stability using the stability function R(z)."
        answer_s = "A method is stable for z=hλ if \\(|R(z)|\\le 1\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # 3) Forward Euler region and step-size rule for real λ<0
    # ------------------------------------------------------------
    if kind == "forward_euler_stability_region":
        z_expr = h * lam
        R = 1 + z
        # for real negative lambda rule: -2 <= z <= 0
        # FIX: Replaced < with \lt to prevent HTML parsing errors
        body = (
            "& \\displaystyle \\text{Forward Euler: } y_{n+1}=y_n+h\\lambda y_n=(1+z)y_n \\\\[6pt]"
            "& \\displaystyle R(z)=1+z \\\\[8pt]"
            "& \\displaystyle \\text{Stability: } |1+z|\\le 1 \\\\[6pt]"
            "& \\displaystyle \\text{If }\\lambda \\lt 0\\text{ is real: } -2\\le z\\le 0 \\Rightarrow 0 \\lt h\\le \\frac{2}{|\\lambda|} \\\\[10pt]"
            "& \\displaystyle \\text{Example: } \\lambda=" + Ltx(lam_val) + ",\\; h=" + Ltx(h_val) + ",\\; z=h\\lambda=" + Ltx(sp.simplify(h_val * lam_val) if simplify else h_val * lam_val) + " \\\\"
        )
        soln_s = card(
            "Forward Euler — Stability Function and Real-Step Rule",
            body,
            "Forward Euler has a bounded stability region, so stiff decay modes (large negative λ) force tiny h for stability.",
        )
        problem_s = "Derive Euler’s stability condition for y' = λy."
        # FIX: Replaced < with \lt in answer string as well for consistency
        answer_s = "Euler: \\(R(z)=1+z\\). Stable when \\(|1+z|\\le 1\\) (for real λ \\lt 0: \\(0 \\lt h\\le 2/|\\lambda|\\))."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # 4) Backward Euler A-stability
    # ------------------------------------------------------------
    if kind == "backward_euler_a_stability":
        # FIX: Replaced < and > with \lt and \gt
        body = (
            "& \\displaystyle \\text{Backward Euler: } y_{n+1}=y_n+h\\lambda y_{n+1} \\\\[4pt]"
            "& \\displaystyle \\Rightarrow y_{n+1}=\\frac{1}{1-z}y_n,\\quad z=h\\lambda \\\\[8pt]"
            "& \\displaystyle R(z)=\\frac{1}{1-z} \\\\[8pt]"
            "& \\displaystyle \\text{If }\\Re(z) \\lt 0\\text{ (left half-plane), then } |R(z)| \\lt 1\\text{ for all }h \\gt 0 \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Backward Euler is A-stable (stable for all }\\Re(z) \\lt 0). \\\\"
        )
        soln_s = card(
            "Backward Euler — A-Stability (Left Half-Plane)",
            body,
            "A-stability is why backward Euler is a go-to for stiff decay problems.",
        )
        problem_s = "Compute R(z) for Backward Euler and state its stability property."
        answer_s = "Backward Euler: \\(R(z)=1/(1-z)\\), A-stable (stable for all \\(\\Re(z) \\lt 0\\))."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # 5) Trapezoidal A-stable but not L-stable
    # ------------------------------------------------------------
    if kind == "trapezoidal_a_stability_not_l_stable":
        # FIX: Replaced < with \lt
        body = (
            "& \\displaystyle \\text{Trapezoidal rule: } y_{n+1}=y_n+\\frac{h\\lambda}{2}(y_n+y_{n+1}) \\\\[4pt]"
            "& \\displaystyle \\Rightarrow y_{n+1}=\\frac{1+z/2}{1-z/2}y_n,\\quad z=h\\lambda \\\\[8pt]"
            "& \\displaystyle R(z)=\\frac{1+z/2}{1-z/2} \\\\[8pt]"
            "& \\displaystyle \\text{A-stable: stable for all }\\Re(z) \\lt 0. \\\\[6pt]"
            "& \\displaystyle \\text{But as } z\\to -\\infty,\\; R(z)\\to -1\\;\\Rightarrow\\;\\text{stiff modes may oscillate (not L-stable).} \\\\"
        )
        soln_s = card(
            "Trapezoidal Rule — A-Stable but Not L-Stable",
            body,
            "Crank–Nicolson/trapezoid is accurate and A-stable, but it does not strongly damp very stiff modes (can cause oscillations).",
        )
        problem_s = "Find R(z) for the trapezoidal rule and explain A-stable vs L-stable behavior."
        answer_s = "Trapezoid: \\(R(z)=(1+z/2)/(1-z/2)\\), A-stable; not L-stable since \\(R(z)\\to -1\\) as \\(z\\to-\\infty\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # 6) RK4 stability function
    # ------------------------------------------------------------
    if kind == "rk4_stability_function":
        # RK4 stability polynomial for y'=λy:
        # R(z) = 1 + z + z^2/2 + z^3/6 + z^4/24
        R = 1 + z + z**2 / 2 + z**3 / 6 + z**4 / 24
        body = (
            "& \\displaystyle \\text{RK4 applied to } y'=\\lambda y \\Rightarrow y_{n+1}=R(z)y_n,\\; z=h\\lambda \\\\[8pt]"
            "& \\displaystyle R(z)=1+z+\\frac{z^2}{2}+\\frac{z^3}{6}+\\frac{z^4}{24} \\quad \\text{(4th-order Taylor of }e^z\\text{)} \\\\[10pt]"
            "& \\displaystyle \\text{RK4 stability region is bounded (explicit method). For large negative }z,\\; |R(z)|\\text{ grows }\\Rightarrow\\text{ needs small }h. \\\\"
        )
        soln_s = card(
            "RK4 — Stability Function (Test Equation)",
            body,
            "Even though RK4 is accurate, it is still explicit and has a bounded stability region—stiff problems can force tiny steps.",
        )
        problem_s = "State RK4’s stability function R(z) for the linear test equation."
        answer_s = "\\(R(z)=1+z+z^2/2+z^3/6+z^4/24\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # 7) Stiffness definition and stiffness ratio (linear systems)
    # ------------------------------------------------------------
    if kind == "stiffness_definition_and_ratio":
        # FIX: Replaced < with \lt
        body = (
            "& \\displaystyle \\text{Stiffness (practical): explicit solvers need tiny }h\\text{ for stability, not accuracy.} \\\\[8pt]"
            "& \\displaystyle \\text{Linear system: } \\mathbf{y}'=A\\mathbf{y}.\\; \\text{Eigenvalues }\\lambda_i\\text{ control decay/oscillation.} \\\\[8pt]"
            "& \\displaystyle \\text{If all }\\Re(\\lambda_i) \\lt 0\\text{ and magnitudes vary a lot, problem is stiff.} \\\\[8pt]"
            "& \\displaystyle \\text{Stiffness ratio (rough): } \\kappa \\approx \\frac{\\max |\\Re(\\lambda_i)|}{\\min |\\Re(\\lambda_i)|} \\\\[6pt]"
            "& \\displaystyle \\text{Large }\\kappa\\text{ (e.g., }" + Ltx(kappa_val) + "\\text{) }\\Rightarrow\\text{ stiffness likely.} \\\\"
        )
        soln_s = card(
            "Stiffness — Definition and Eigenvalue Ratio (Linearized View)",
            body,
            "Stiffness is about stability constraints dominating the step size. Implicit A-stable/L-stable methods are preferred when κ is large.",
        )
        problem_s = "Explain stiffness and give the stiffness ratio idea for linear systems."
        answer_s = "Stiff if explicit steps are forced tiny for stability; for \\(y'=Ay\\), \\(\\kappa\\approx\\max|\\Re\\lambda|/\\min|\\Re\\lambda|\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # 8) Stiffness detection heuristics (app-friendly)
    # ------------------------------------------------------------
    if kind == "stiffness_detection_heuristics":
        body = (
            "& \\displaystyle \\text{Heuristics to detect stiffness (without full eigen-analysis):} \\\\[6pt]"
            "& \\displaystyle \\text{(1) Adaptive RK45 step-size collapse: }h\\text{ shrinks drastically while solution seems smooth.} \\\\[4pt]"
            "& \\displaystyle \\text{(2) Many rejected steps in explicit solvers.} \\\\[4pt]"
            "& \\displaystyle \\text{(3) Large }\\left\\|\\frac{\\partial f}{\\partial y}\\right\\|\\text{ (Jacobian norm) in some region.} \\\\[4pt]"
            "& \\displaystyle \\text{(4) For scalar: estimate dominant }\\lambda\\approx \\frac{\\partial f}{\\partial y},\\;\\text{then check explicit stability }|R(h\\lambda)|. \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{If stiffness suspected: switch to implicit (Backward Euler / BDF / implicit trapezoid).} \\\\"
        )
        soln_s = card(
            "Stiffness Detection — Practical Heuristics for a Solver",
            body,
            "These are the decision rules your app can cite when choosing between explicit RK and implicit stiff solvers.",
        )
        problem_s = "List practical stiffness detection heuristics suitable for a solver/app."
        answer_s = "Step collapse, frequent rejections, large Jacobian norm, estimated λ causing explicit instability ⇒ use implicit method."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # 9) PDE CFL rules (tie-in to numerical PDE bucket)
    # ------------------------------------------------------------
    if kind == "pde_cfl_rules":
        body = (
            "& \\displaystyle \\text{Explicit PDE schemes have CFL-style restrictions:} \\\\[8pt]"
            "& \\displaystyle \\text{Heat FTCS (1D): } r=\\frac{k\\Delta t}{(\\Delta x)^2}\\le \\frac12 \\;\\Rightarrow\\; \\Delta t\\le \\frac{(\\Delta x)^2}{2k} \\\\[8pt]"
            "& \\displaystyle \\text{Wave (1D centered): } \\lambda=\\frac{c\\Delta t}{\\Delta x}\\le 1 \\;\\Rightarrow\\; \\Delta t\\le \\frac{\\Delta x}{c} \\\\[10pt]"
            "& \\displaystyle \\text{These are amplification-factor conditions (von Neumann stability) specialized to common stencils.} \\\\"
        )
        soln_s = card(
            "PDE Stability Rules — CFL Conditions (Heat/Wave)",
            body,
            "This ties stability analysis to grid choices (Δx, Δt). Your app can use these inequalities to auto-suggest safe timesteps.",
        )
        problem_s = "State standard CFL/stability conditions for explicit heat and wave schemes."
        answer_s = "Heat FTCS: \\(k\\Delta t/(\\Delta x)^2\\le 1/2\\). Wave: \\(c\\Delta t/\\Delta x\\le 1\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError(
        "Unknown kind={!r}. Use one of: "
        "'linear_test_equation_overview','absolute_stability_definition',"
        "'forward_euler_stability_region','backward_euler_a_stability',"
        "'trapezoidal_a_stability_not_l_stable','rk4_stability_function',"
        "'stiffness_definition_and_ratio','stiffness_detection_heuristics','pde_cfl_rules'.".format(kind)
    )