# ============================================================
# File: _solve_implicit_nonlinear_solve_newton_engine.py
# Topic: Implicit nonlinear solve machinery (Newton / quasi-Newton)
#
# This file provides an "explain + template" solver that:
#   - builds the residual equation F(y)=0 for common implicit steps
#   - shows Newton update (scalar) and Newton-Krylov style (vector) forms
#   - shows Jacobian strategies: analytic, finite-diff, quasi-Newton (Broyden)
#   - shows damping / line search + stopping criteria
#   - shows linear-solve substep (tri-diagonal / sparse) in Newton
#
# It is NOT a numeric engine here (no actual iteration over floats),
# but it produces a canonical step-by-step derivation and pseudocode,
# consistent with your "problem template" outputs.
#
# Output:
#   { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
#
# Notes:
#   - Avoid f-strings to prevent LaTeX brace issues.
#   - Uses expr.doit() method ONLY if needed (not sp.doit).
# ============================================================

from __future__ import annotations

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

import sympy as sp


ImplicitNewtonKind = Literal[
    "implicit_ode_backward_euler_scalar",
    "implicit_ode_backward_euler_vector",
    "implicit_ode_trapezoid_scalar",
    "implicit_pde_step_template",
    "newton_scalar_core",
    "newton_vector_core",
    "finite_difference_jacobian_template",
    "broyden_quasi_newton_template",
    "damped_newton_line_search_template",
    "stopping_criteria_template",
]


def solve_implicit_nonlinear_solve_newton_engine(
    *,
    kind: ImplicitNewtonKind,
    t_symbol: str = "t",
    y_symbol: str = "y",
    h_symbol: str = "h",
    f_symbol: str = "f",
    tol_symbol: str = "tol",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    # Basic symbols
    t = sp.Symbol(t_symbol, real=True)
    h = sp.Symbol(h_symbol, positive=True, real=True)
    tol = sp.Symbol(tol_symbol, positive=True, real=True)

    # Scalar unknown at new time
    yn = sp.Symbol("y_n", real=True)
    yn1 = sp.Symbol("y_{n+1}", real=True)
    tn = sp.Symbol("t_n", real=True)
    tn1 = sp.Symbol("t_{n+1}", real=True)

    # Vector symbols
    Y = sp.MatrixSymbol("Y", 1, 1)   # placeholder, used for latex only
    Yn = sp.MatrixSymbol("Y_n", 1, 1)
    Yn1 = sp.MatrixSymbol("Y_{n+1}", 1, 1)

    # Functions f(t,y) and residual F(y)
    f = sp.Function(f_symbol)
    Ffun = sp.Function("F")

    base_locals = {
        t_symbol: t,
        h_symbol: h,
        tol_symbol: tol,
        "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>"
        )

    # =========================================================
    # 1) Implicit ODE: Backward Euler (scalar)
    # =========================================================
    if kind == "implicit_ode_backward_euler_scalar":
        body = (
            "& \\displaystyle y' = f(t,y) \\\\[6pt]"
            "& \\displaystyle \\text{Backward Euler step:}\\quad y_{n+1}=y_n+h\\,f(t_{n+1},y_{n+1}) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Define residual }F(y)=y-y_n-h f(t_{n+1},y) \\\\[4pt]"
            "& \\displaystyle \\Rightarrow \\text{Solve }F(y_{n+1})=0 \\\\[10pt]"
            "& \\displaystyle \\text{Newton iteration: } y^{(k+1)}=y^{(k)}-\\frac{F(y^{(k)})}{F'(y^{(k)})} \\\\[6pt]"
            "& \\displaystyle F'(y)=1-h\\,\\frac{\\partial f}{\\partial y}(t_{n+1},y) \\\\[6pt]"
            "& \\displaystyle \\text{Good initial guess (predictor): } y^{(0)}=y_n+h f(t_n,y_n) \\\\"
        )
        soln_s = card(
            "Implicit ODE Step — Backward Euler Residual + Newton Solve (Scalar)",
            body,
            "Implicit methods convert each step into a root-finding problem. Newton is the standard engine behind stiff solvers.",
        )
        problem_s = "Form the nonlinear equation and Newton iteration for Backward Euler."
        answer_s = "Solve \\(F(y)=y-y_n-h f(t_{n+1},y)=0\\) by Newton: \\(y^{(k+1)}=y^{(k)}-F/F'\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 2) Implicit ODE: Backward Euler (vector system)
    # =========================================================
    if kind == "implicit_ode_backward_euler_vector":
        body = (
            "& \\displaystyle \\mathbf{y}'=\\mathbf{f}(t,\\mathbf{y}) \\\\[6pt]"
            "& \\displaystyle \\text{Backward Euler: } \\mathbf{y}_{n+1}=\\mathbf{y}_n+h\\,\\mathbf{f}(t_{n+1},\\mathbf{y}_{n+1}) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\mathbf{F}(\\mathbf{y})=\\mathbf{y}-\\mathbf{y}_n-h\\,\\mathbf{f}(t_{n+1},\\mathbf{y})=\\mathbf{0} \\\\[8pt]"
            "& \\displaystyle \\text{Vector Newton: solve } J(\\mathbf{y}^{(k)})\\,\\Delta^{(k)}=-\\mathbf{F}(\\mathbf{y}^{(k)}) \\\\[4pt]"
            "& \\displaystyle \\text{Update: } \\mathbf{y}^{(k+1)}=\\mathbf{y}^{(k)}+\\Delta^{(k)} \\\\[10pt]"
            "& \\displaystyle \\text{Jacobian: } J=\\frac{\\partial \\mathbf{F}}{\\partial \\mathbf{y}} = I - h\\,\\frac{\\partial \\mathbf{f}}{\\partial \\mathbf{y}}(t_{n+1},\\mathbf{y}) \\\\"
        )
        soln_s = card(
            "Implicit ODE Step — Backward Euler + Vector Newton (Jacobian Linear Solve)",
            body,
            "In practice you never invert the Jacobian; you solve a linear system each Newton iteration.",
        )
        problem_s = "Write the vector Newton form for an implicit Backward Euler step."
        answer_s = "Solve \\(\\mathbf{F}(\\mathbf{y})=\\mathbf{0}\\) with \\(J\\Delta=-\\mathbf{F}\\), \\(\\mathbf{y}\\leftarrow\\mathbf{y}+\\Delta\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 3) Implicit ODE: Trapezoid / Crank–Nicolson (scalar)
    # =========================================================
    if kind == "implicit_ode_trapezoid_scalar":
        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) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow F(y)=y-y_n-\\frac{h}{2}\\Big(f(t_n,y_n)+f(t_{n+1},y)\\Big)=0 \\\\[10pt]"
            "& \\displaystyle \\text{Newton: } y^{(k+1)}=y^{(k)}-\\frac{F(y^{(k)})}{F'(y^{(k)})} \\\\[6pt]"
            "& \\displaystyle F'(y)=1-\\frac{h}{2}\\,\\frac{\\partial f}{\\partial y}(t_{n+1},y) \\\\[6pt]"
            "& \\displaystyle \\text{Predictor for }y^{(0)}:\\; y_n+h f(t_n,y_n) \\text{ (or Heun predictor)} \\\\"
        )
        soln_s = card(
            "Implicit ODE Step — Trapezoid/Crank–Nicolson Residual + Newton (Scalar)",
            body,
            "Crank–Nicolson is second-order but implicit. A predictor often improves Newton robustness.",
        )
        problem_s = "Form the nonlinear equation for implicit trapezoid and the Newton derivative."
        answer_s = "Solve \\(F(y)=y-y_n-\\tfrac{h}{2}(f(t_n,y_n)+f(t_{n+1},y))=0\\) via Newton."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 4) Implicit PDE step (template)
    # =========================================================
    if kind == "implicit_pde_step_template":
        body = (
            "& \\displaystyle \\text{After spatial discretization of a PDE, you get a system } \\mathbf{U}^{n+1}\\text{ to solve.} \\\\[6pt]"
            "& \\displaystyle \\text{Generic implicit step: } \\mathbf{F}(\\mathbf{U}^{n+1})=\\mathbf{0} \\\\[8pt]"
            "& \\displaystyle \\text{If PDE is linear: } \\mathbf{F}(\\mathbf{U})=A\\mathbf{U}-\\mathbf{b} \\Rightarrow \\text{solve linear system.} \\\\[8pt]"
            "& \\displaystyle \\text{If PDE is nonlinear: use Newton iterations: } J(\\mathbf{U}^{(k)})\\Delta=-\\mathbf{F}(\\mathbf{U}^{(k)}) \\\\[4pt]"
            "& \\displaystyle \\mathbf{U}^{(k+1)}=\\mathbf{U}^{(k)}+\\Delta \\\\[8pt]"
            "& \\displaystyle \\text{Linear solve choices: tridiagonal (1D), sparse/direct, or iterative (CG/GMRES).} \\\\"
        )
        soln_s = card(
            "Implicit PDE Time-Step — Nonlinear Residual + Linear Solve Substep (Template)",
            body,
            "This is the reusable engine behind BTCS / Crank–Nicolson for PDEs, especially when the PDE is nonlinear.",
        )
        problem_s = "Explain how implicit PDE timesteps reduce to nonlinear (or linear) algebraic solves."
        answer_s = "Implicit timestep ⇒ solve \\(\\mathbf{F}(\\mathbf{U}^{n+1})=0\\) (linear or Newton with Jacobian solves)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 5) Newton core (scalar)
    # =========================================================
    if kind == "newton_scalar_core":
        body = (
            "& \\displaystyle \\text{Goal: solve } F(y)=0 \\\\[8pt]"
            "& \\displaystyle \\text{Newton update: } y^{(k+1)} = y^{(k)} - \\frac{F(y^{(k)})}{F'(y^{(k)})} \\\\[10pt]"
            "& \\displaystyle \\text{Stopping: } |F(y^{(k)})|<\\text{tol} \\;\\text{or}\\; |y^{(k+1)}-y^{(k)}|<\\text{tol}(1+|y^{(k)}|) \\\\"
        )
        soln_s = card(
            "Nonlinear Solve Engine — Newton’s Method (Scalar Core)",
            body,
            "This is the root-solver that implicit time-steppers call every timestep.",
        )
        problem_s = "State Newton’s method for a scalar nonlinear equation and common stopping criteria."
        answer_s = "\\(y^{(k+1)}=y^{(k)}-F/F'\\) with residual/step stopping tests."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 6) Newton core (vector)
    # =========================================================
    if kind == "newton_vector_core":
        body = (
            "& \\displaystyle \\text{Goal: solve } \\mathbf{F}(\\mathbf{y})=\\mathbf{0} \\\\[8pt]"
            "& \\displaystyle \\text{Jacobian: } J_{ij}=\\frac{\\partial F_i}{\\partial y_j} \\\\[6pt]"
            "& \\displaystyle \\text{Newton step: solve } J(\\mathbf{y}^{(k)})\\Delta^{(k)}=-\\mathbf{F}(\\mathbf{y}^{(k)}) \\\\[4pt]"
            "& \\displaystyle \\text{Update: } \\mathbf{y}^{(k+1)}=\\mathbf{y}^{(k)}+\\Delta^{(k)} \\\\[10pt]"
            "& \\displaystyle \\text{Never compute }J^{-1}.\\;\\text{Always solve the linear system.} \\\\"
        )
        soln_s = card(
            "Nonlinear Solve Engine — Newton’s Method (Vector Core)",
            body,
            "Vector Newton = repeated linear solves. This is why sparse/tridiagonal solvers matter for PDE grids.",
        )
        problem_s = "State Newton’s method for systems and the Jacobian linear solve form."
        answer_s = "Solve \\(J\\Delta=-\\mathbf{F}\\), then \\(\\mathbf{y}\\leftarrow\\mathbf{y}+\\Delta\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 7) Finite-difference Jacobian template
    # =========================================================
    if kind == "finite_difference_jacobian_template":
        eps = sp.Symbol("\\varepsilon", positive=True, real=True)
        body = (
            "& \\displaystyle \\text{If analytic }J\\text{ is unavailable, approximate }J\\text{ by finite differences.} \\\\[8pt]"
            "& \\displaystyle \\frac{\\partial F_i}{\\partial y_j}\\approx \\frac{F_i(\\mathbf{y}+\\varepsilon\\mathbf{e}_j)-F_i(\\mathbf{y})}{\\varepsilon} \\\\[10pt]"
            "& \\displaystyle \\text{Pros: universal.}\\;\\text{Cons: costs }O(n)\\text{ residual evaluations per Jacobian build.} \\\\"
        )
        soln_s = card(
            "Jacobian Strategy — Finite-Difference Jacobian (Template)",
            body,
            "Good fallback for general problems, but expensive for large systems.",
        )
        problem_s = "Write the finite-difference approximation of the Jacobian for Newton."
        answer_s = "\\(J_{ij}\\approx (F_i(y+\\varepsilon e_j)-F_i(y))/\\varepsilon\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 8) Broyden quasi-Newton template
    # =========================================================
    if kind == "broyden_quasi_newton_template":
        body = (
            "& \\displaystyle \\text{Quasi-Newton (Broyden): avoid recomputing full Jacobian each iteration.} \\\\[8pt]"
            "& \\displaystyle \\text{Maintain an approximation }B^{(k)}\\approx J(\\mathbf{y}^{(k)}) \\\\[6pt]"
            "& \\displaystyle \\text{Solve }B^{(k)}\\Delta^{(k)}=-\\mathbf{F}(\\mathbf{y}^{(k)}) \\\\[6pt]"
            "& \\displaystyle \\mathbf{y}^{(k+1)}=\\mathbf{y}^{(k)}+\\Delta^{(k)} \\\\[10pt]"
            "& \\displaystyle \\text{Update }B^{(k)}\\text{ using secant condition } B^{(k+1)}s^{(k)}=q^{(k)} \\\\[4pt]"
            "& \\displaystyle s^{(k)}=\\mathbf{y}^{(k+1)}-\\mathbf{y}^{(k)},\\quad q^{(k)}=\\mathbf{F}(\\mathbf{y}^{(k+1)})-\\mathbf{F}(\\mathbf{y}^{(k)}) \\\\"
        )
        soln_s = card(
            "Quasi-Newton — Broyden Update (Template)",
            body,
            "Broyden trades accuracy for fewer Jacobian builds; useful when residual evaluation is cheap but Jacobian is expensive.",
        )
        problem_s = "Describe the quasi-Newton (Broyden) idea and the secant update condition."
        answer_s = "Maintain \\(B\\approx J\\), solve \\(B\\Delta=-F\\), and update \\(B\\) to satisfy \\(B s=q\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 9) Damped Newton / line search template
    # =========================================================
    if kind == "damped_newton_line_search_template":
        alpha = sp.Symbol("\\alpha", positive=True, real=True)
        body = (
            "& \\displaystyle \\text{If full Newton steps overshoot, use damping:} \\\\[6pt]"
            "& \\displaystyle \\mathbf{y}^{(k+1)}=\\mathbf{y}^{(k)}+\\alpha\\,\\Delta^{(k)},\\quad 0<\\alpha\\le 1 \\\\[10pt]"
            "& \\displaystyle \\text{Line search idea: reduce }\\alpha\\text{ until } \\|\\mathbf{F}(\\mathbf{y}^{(k+1)})\\| < \\|\\mathbf{F}(\\mathbf{y}^{(k)})\\| \\\\"
        )
        soln_s = card(
            "Newton Safeguard — Damped Newton / Line Search (Template)",
            body,
            "Damping makes Newton far more robust for difficult nonlinearities and large timesteps.",
        )
        problem_s = "Write the damped Newton update and the purpose of line search."
        answer_s = "\\(y^{k+1}=y^k+\\alpha\\Delta^k\\) with \\(0<\\alpha\\le 1\\), reduce \\(\\alpha\\) if residual increases."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 10) Stopping criteria template
    # =========================================================
    if kind == "stopping_criteria_template":
        body = (
            "& \\displaystyle \\text{Common stopping tests:} \\\\[6pt]"
            "& \\displaystyle \\text{(A) Residual: } \\|\\mathbf{F}(\\mathbf{y}^{(k)})\\|<\\text{tol} \\\\[6pt]"
            "& \\displaystyle \\text{(B) Step: } \\|\\Delta^{(k)}\\|<\\text{tol}\\,(1+\\|\\mathbf{y}^{(k)}\\|) \\\\[10pt]"
            "& \\displaystyle \\text{Safeguards: max-iter, damping/line-search, or reduce timestep }h\\text{ if failing.} \\\\"
        )
        soln_s = card(
            "Newton Convergence — Stopping Criteria + Fallback Policies (Template)",
            body,
            "For implicit time stepping, failure to converge is usually handled by reducing the timestep and retrying.",
        )
        problem_s = "List standard Newton stopping criteria and common fallback policies."
        answer_s = "Stop when residual small or step small; otherwise damp Newton / limit iterations / reduce timestep."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError(
        "Unknown kind={!r}. Use one of: "
        "'implicit_ode_backward_euler_scalar','implicit_ode_backward_euler_vector',"
        "'implicit_ode_trapezoid_scalar','implicit_pde_step_template',"
        "'newton_scalar_core','newton_vector_core',"
        "'finite_difference_jacobian_template','broyden_quasi_newton_template',"
        "'damped_newton_line_search_template','stopping_criteria_template'.".format(kind)
    )
