# ============================================================
# File: _solve_numerical_pde_finite_difference_methods.py
# Topic: Numerical PDE Methods (Finite Differences + Stability/CFL)
#
# Covers "textbook + app-friendly" templates for:
#   1) Heat equation (1D): FTCS / BTCS / Crank–Nicolson + stability
#   2) Wave equation (1D): centered explicit scheme + CFL + first step
#   3) Laplace/Poisson (2D steady): 5-point stencil + Jacobi/Gauss-Seidel/SOR notes
#
# We do NOT attempt to run heavy numerics here; we generate:
#   - the scheme (stencil / update rule / linear system)
#   - the stability ratio and CFL checks
#   - pseudocode-style steps
#
# Output:
#   { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
#
# Notes:
#   - Avoid f-strings (to prevent brace issues with LaTeX).
#   - SymPy used for parsing/simplify/latex only.
# ============================================================

from __future__ import annotations

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

import sympy as sp


NumericalPDEKind = Literal[
    # Heat (1D)
    "heat_ftcs",
    "heat_btcs",
    "heat_crank_nicolson",
    # Wave (1D)
    "wave_centered",
    # Elliptic (2D)
    "laplace_2d_stencil_jacobi",
    "poisson_2d_stencil_jacobi",
]


def solve_numerical_pde_finite_difference_methods(
    *,
    kind: NumericalPDEKind,

    # Common domain/grid parameters
    L_: Union[int, float, str, sp.Expr] = 1,
    H_: Union[int, float, str, sp.Expr] = 1,   # for 2D
    Nx_: int = 10,
    Ny_: int = 10,                             # for 2D
    dt_: Union[int, float, str, sp.Expr] = sp.Rational(1, 100),
    T_: Union[int, float, str, sp.Expr] = 1,

    # PDE parameters
    k_: Union[int, float, str, sp.Expr] = 1,   # heat diffusivity
    c_: Union[int, float, str, sp.Expr] = 1,   # wave speed

    # For Poisson source term symbolically (2D)
    f_of_xy: Union[str, sp.Expr] = "0",

    # Symbols
    x_symbol: str = "x",
    y_symbol: str = "y",
    t_symbol: str = "t",
    u_symbol: str = "u",

    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    x = sp.Symbol(x_symbol, real=True)
    y = sp.Symbol(y_symbol, real=True)
    t = sp.Symbol(t_symbol, real=True)
    u = sp.Function(u_symbol)

    base_locals = {
        x_symbol: x,
        y_symbol: y,
        t_symbol: t,
        u_symbol: u,
        "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 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>"
        )

    L = P(L_)
    H = P(H_)
    dt = P(dt_)
    T = P(T_)
    k = P(k_)
    c = P(c_)

    if simplify:
        L = sp.simplify(L)
        H = sp.simplify(H)
        dt = sp.simplify(dt)
        T = sp.simplify(T)
        k = sp.simplify(k)
        c = sp.simplify(c)

    Nx = int(Nx_)
    Ny = int(Ny_)

    if Nx < 2:
        Nx = 2
    if Ny < 2:
        Ny = 2

    dx = L / Nx
    dy = H / Ny

    if simplify:
        dx = sp.simplify(dx)
        dy = sp.simplify(dy)

    # Stability ratios
    r_heat = k * dt / (dx**2)
    lam_wave = c * dt / dx

    if simplify:
        r_heat = sp.simplify(r_heat)
        lam_wave = sp.simplify(lam_wave)

    # ---------- helpers for stability text ----------
    def heat_stability_line() -> str:
        return (
            "& \\displaystyle r=\\frac{k\\,\\Delta t}{(\\Delta x)^2}="
            + Ltx(r_heat)
            + ",\\quad \\text{FTCS stability (1D): } r\\le \\frac12 \\\\"
        )

    def wave_cfl_line() -> str:
        return (
            "& \\displaystyle \\lambda=\\frac{c\\,\\Delta t}{\\Delta x}="
            + Ltx(lam_wave)
            + ",\\quad \\text{CFL stability (1D wave): } \\lambda\\le 1 \\\\"
        )

    # =========================================================
    # HEAT: FTCS
    # =========================================================
    if kind == "heat_ftcs":
        body = (
            "& \\displaystyle u_t = k u_{xx},\\quad 0<x<" + Ltx(L) + " \\\\"
            "& \\displaystyle x_i=i\\Delta x,\\; t^n=n\\Delta t,\\; u_i^n\\approx u(x_i,t^n) \\\\"
            "& \\displaystyle \\Delta x=" + Ltx(dx) + ",\\; \\Delta t=" + Ltx(dt) + ",\\; N_x=" + str(Nx) + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow u_t\\approx \\frac{u_i^{n+1}-u_i^n}{\\Delta t},\\quad u_{xx}\\approx \\frac{u_{i+1}^n-2u_i^n+u_{i-1}^n}{(\\Delta x)^2} \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\frac{u_i^{n+1}-u_i^n}{\\Delta t}=k\\frac{u_{i+1}^n-2u_i^n+u_{i-1}^n}{(\\Delta x)^2} \\\\[6pt]"
            "& \\displaystyle \\Rightarrow u_i^{n+1}=u_i^n + r\\,(u_{i+1}^n-2u_i^n+u_{i-1}^n) \\\\[6pt]"
            + heat_stability_line() +
            "& \\displaystyle \\Rightarrow \\text{Pseudocode (Dirichlet BC example):} \\\\"
            "& \\displaystyle \\text{for }n=0..N_t-1:\\;\\; u_0^{n}=g_0(t^n),\\; u_{N_x}^{n}=g_L(t^n) \\\\"
            "& \\displaystyle \\quad \\text{for }i=1..N_x-1:\\; u_i^{n+1}=u_i^n+r(u_{i+1}^n-2u_i^n+u_{i-1}^n) \\\\"
        )
        soln_s = card(
            "Numerical PDE — Heat Equation (FTCS Explicit) + Stability",
            body,
            "FTCS is simple but conditionally stable: choose Δt small enough so r ≤ 1/2 in 1D.",
        )
        problem_s = "Discretize the 1D heat equation using FTCS and state the stability condition."
        answer_s = (
            "\\(u_i^{n+1}=u_i^n+r(u_{i+1}^n-2u_i^n+u_{i-1}^n),\\; r=\\frac{k\\Delta t}{(\\Delta x)^2}\\le\\frac12\\)."
        )
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # HEAT: BTCS (Implicit)
    # =========================================================
    if kind == "heat_btcs":
        body = (
            "& \\displaystyle u_t = k u_{xx},\\quad 0<x<" + Ltx(L) + " \\\\"
            "& \\displaystyle \\Delta x=" + Ltx(dx) + ",\\; \\Delta t=" + Ltx(dt) + ",\\; r=" + Ltx(r_heat) + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\frac{u_i^{n+1}-u_i^n}{\\Delta t}=k\\frac{u_{i+1}^{n+1}-2u_i^{n+1}+u_{i-1}^{n+1}}{(\\Delta x)^2} \\\\[6pt]"
            "& \\displaystyle \\Rightarrow -r\\,u_{i-1}^{n+1}+(1+2r)u_i^{n+1}-r\\,u_{i+1}^{n+1}=u_i^n \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{This is a tridiagonal system each step: } A\\,\\mathbf{u}^{n+1}=\\mathbf{u}^n+\\text{(BC terms)} \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Stability: BTCS is unconditionally stable (no r\\le 1/2 restriction).} \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Thomas algorithm (outline): forward elimination + back substitution.} \\\\"
        )
        soln_s = card(
            "Numerical PDE — Heat Equation (BTCS Implicit) + Tridiagonal Solve",
            body,
            "BTCS costs a linear solve per step but is stable for any Δt (for the linear heat equation).",
        )
        problem_s = "Discretize the 1D heat equation using BTCS and write the tridiagonal system."
        answer_s = "\\(-r u_{i-1}^{n+1}+(1+2r)u_i^{n+1}-r u_{i+1}^{n+1}=u_i^n\\) (tridiagonal), BTCS stable for any Δt."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # HEAT: Crank–Nicolson
    # =========================================================
    if kind == "heat_crank_nicolson":
        body = (
            "& \\displaystyle u_t = k u_{xx},\\quad 0<x<" + Ltx(L) + " \\\\"
            "& \\displaystyle \\Delta x=" + Ltx(dx) + ",\\; \\Delta t=" + Ltx(dt) + ",\\; r=" + Ltx(r_heat) + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\frac{u_i^{n+1}-u_i^n}{\\Delta t}"
            "=\\frac{k}{2}\\left[\\frac{u_{i+1}^{n}-2u_i^{n}+u_{i-1}^{n}}{(\\Delta x)^2}"
            "+\\frac{u_{i+1}^{n+1}-2u_i^{n+1}+u_{i-1}^{n+1}}{(\\Delta x)^2}\\right] \\\\[8pt]"
            "& \\displaystyle \\Rightarrow -\\frac{r}{2}u_{i-1}^{n+1}+(1+r)u_i^{n+1}-\\frac{r}{2}u_{i+1}^{n+1}"
            "=\\frac{r}{2}u_{i-1}^{n}+(1-r)u_i^{n}+\\frac{r}{2}u_{i+1}^{n} \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Tridiagonal system each step; 2nd-order accurate in time.} \\\\"
        )
        soln_s = card(
            "Numerical PDE — Heat Equation (Crank–Nicolson) + Tridiagonal Solve",
            body,
            "Crank–Nicolson is typically very accurate; with sharp initial data, large Δt can introduce mild oscillations.",
        )
        problem_s = "Write the Crank–Nicolson scheme for the 1D heat equation."
        answer_s = "\\(-\\tfrac r2 u_{i-1}^{n+1}+(1+r)u_i^{n+1}-\\tfrac r2 u_{i+1}^{n+1}=\\tfrac r2 u_{i-1}^{n}+(1-r)u_i^{n}+\\tfrac r2 u_{i+1}^{n}\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # WAVE: Centered explicit (2nd order in time)
    # =========================================================
    if kind == "wave_centered":
        body = (
            "& \\displaystyle u_{tt}=c^2 u_{xx},\\quad 0<x<" + Ltx(L) + " \\\\"
            "& \\displaystyle \\Delta x=" + Ltx(dx) + ",\\; \\Delta t=" + Ltx(dt) + ",\\; \\lambda=" + Ltx(lam_wave) + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow u_{tt}\\approx \\frac{u_i^{n+1}-2u_i^n+u_i^{n-1}}{(\\Delta t)^2},\\quad"
            "u_{xx}\\approx \\frac{u_{i+1}^n-2u_i^n+u_{i-1}^n}{(\\Delta x)^2} \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\frac{u_i^{n+1}-2u_i^n+u_i^{n-1}}{(\\Delta t)^2}"
            "=c^2\\frac{u_{i+1}^n-2u_i^n+u_{i-1}^n}{(\\Delta x)^2} \\\\[6pt]"
            "& \\displaystyle \\Rightarrow u_i^{n+1}=2u_i^n-u_i^{n-1}+\\lambda^2(u_{i+1}^n-2u_i^n+u_{i-1}^n) \\\\[6pt]"
            + wave_cfl_line() +
            "& \\displaystyle \\Rightarrow \\text{First step (needs initial velocity }u_t(x,0)=g(x)\\text{):} \\\\"
            "& \\displaystyle u_i^{1}=u_i^{0}+\\Delta t\\,g(x_i)+\\frac{\\lambda^2}{2}(u_{i+1}^0-2u_i^0+u_{i-1}^0) \\\\"
        )
        soln_s = card(
            "Numerical PDE — Wave Equation (Centered Explicit) + CFL",
            body,
            "This scheme is stable if λ ≤ 1 (1D). You must compute u^1 using initial velocity.",
        )
        problem_s = "Discretize the 1D wave equation using the centered explicit scheme and state the CFL condition."
        answer_s = "\\(u_i^{n+1}=2u_i^n-u_i^{n-1}+\\lambda^2(u_{i+1}^n-2u_i^n+u_{i-1}^n),\\; \\lambda=\\frac{c\\Delta t}{\\Delta x}\\le 1\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # LAPLACE: 2D stencil + Jacobi iteration
    # =========================================================
    if kind == "laplace_2d_stencil_jacobi":
        body = (
            "& \\displaystyle u_{xx}+u_{yy}=0,\\quad (x,y)\\in(0," + Ltx(L) + ")\\times(0," + Ltx(H) + ") \\\\"
            "& \\displaystyle \\Delta x=" + Ltx(dx) + ",\\; \\Delta y=" + Ltx(dy) + ",\\; N_x=" + str(Nx) + ",\\; N_y=" + str(Ny) + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{5-point stencil (assuming }\\Delta x=\\Delta y\\text{ for simplicity):} \\\\"
            "& \\displaystyle u_{i+1,j}+u_{i-1,j}+u_{i,j+1}+u_{i,j-1}-4u_{i,j}=0 \\\\[6pt]"
            "& \\displaystyle \\Rightarrow u_{i,j}=\\frac{1}{4}(u_{i+1,j}+u_{i-1,j}+u_{i,j+1}+u_{i,j-1}) \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Jacobi iteration: } u_{i,j}^{(new)}=\\frac14( u_{i+1,j}^{(old)}+u_{i-1,j}^{(old)}+u_{i,j+1}^{(old)}+u_{i,j-1}^{(old)} ) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Stop when max change < tolerance.} \\\\"
        )
        soln_s = card(
            "Numerical PDE — 2D Laplace (5-Point Stencil) + Jacobi Iteration",
            body,
            "For faster convergence use Gauss–Seidel or SOR; boundaries are imposed directly each iteration (Dirichlet).",
        )
        problem_s = "Write the 5-point finite-difference stencil for Laplace’s equation and the Jacobi iteration."
        answer_s = "\\(u_{i,j}=\\tfrac14(u_{i+1,j}+u_{i-1,j}+u_{i,j+1}+u_{i,j-1})\\) (interior), iterate until convergence."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # POISSON: 2D stencil + Jacobi iteration
    # =========================================================
    if kind == "poisson_2d_stencil_jacobi":
        f = P(f_of_xy) if isinstance(f_of_xy, str) else sp.sympify(f_of_xy, locals=base_locals)
        if simplify:
            f = sp.simplify(f)

        body = (
            "& \\displaystyle u_{xx}+u_{yy}=f(x,y),\\quad (x,y)\\in(0," + Ltx(L) + ")\\times(0," + Ltx(H) + ") \\\\"
            "& \\displaystyle \\Delta x=" + Ltx(dx) + ",\\; \\Delta y=" + Ltx(dy) + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Assume }\\Delta x=\\Delta y\\text{ (common case) and define } f_{i,j}=f(x_i,y_j). \\\\"
            "& \\displaystyle \\Rightarrow u_{i+1,j}+u_{i-1,j}+u_{i,j+1}+u_{i,j-1}-4u_{i,j}=(\\Delta x)^2 f_{i,j} \\\\[8pt]"
            "& \\displaystyle \\Rightarrow u_{i,j}=\\frac{1}{4}\\left(u_{i+1,j}+u_{i-1,j}+u_{i,j+1}+u_{i,j-1}-(\\Delta x)^2 f_{i,j}\\right) \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Jacobi iteration uses RHS with old values; boundary values fixed each sweep.} \\\\"
            "& \\displaystyle \\Rightarrow \\text{Source example used: } f(x,y)=" + Ltx(f) + " \\\\"
        )
        soln_s = card(
            "Numerical PDE — 2D Poisson (5-Point Stencil) + Jacobi Iteration",
            body,
            "This is the standard steady-state discretization; for performance use Gauss–Seidel/SOR or sparse linear solvers.",
        )
        problem_s = "Write the 5-point finite-difference stencil for Poisson’s equation and the Jacobi update."
        answer_s = "\\(u_{i,j}=\\tfrac14(u_{i+1,j}+u_{i-1,j}+u_{i,j+1}+u_{i,j-1}-(\\Delta x)^2 f_{i,j})\\) (for Δx=Δy)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError(
        "Unknown kind={!r}. Use "
        "'heat_ftcs', 'heat_btcs', 'heat_crank_nicolson', "
        "'wave_centered', 'laplace_2d_stencil_jacobi', 'poisson_2d_stencil_jacobi'.".format(kind)
    )
