# ============================================================
# File: _solve_pde_parabolic.py
# Topic: Parabolic PDEs (Heat/Diffusion) — App-Friendly Canonical Cases
#
# Parabolic PDE hallmark (2 variables): one time-like + one space-like second derivative
#   Example: Heat equation  u_t = k u_xx   (parabolic)
#
# We implement a small solver for a few classic parabolic PDE cases:
#
# (1) Heat equation on a finite rod (0<x<L) with zero end temperatures:
#       u_t = k u_xx,  0<x<L, t>0
#       u(0,t)=0, u(L,t)=0
#       u(x,0) = sin(nπx/L)        (single sine eigenmode)
#     Solution:
#       u(x,t)=sin(nπx/L) * exp(-k (nπ/L)^2 t)
#
# (2) Heat equation on a finite rod with fixed nonzero end temperatures via steady-state + transient:
#       u_t = k u_xx, 0<x<L
#       u(0,t)=0, u(L,t)=U0  (constants)
#       u(x,0)=0
#     Let steady state: u_s(x)=U0 * x/L
#     Let v=u-u_s => v_t = k v_xx with homogeneous BC and IC v(x,0)=-u_s(x)
#     (We return the decomposition and the Fourier-sine series form for v.)
#
# Notes:
#   - We intentionally keep it closed-form (mode solution) or series-form (standard).
#   - HTML output uses \\(\\displaystyle ...\\) and aligned steps.
# ============================================================

from __future__ import annotations

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

import sympy as sp


ParabolicPdeKind = Literal[
    "heat_1d_sine_mode",
    "heat_1d_fixed_ends_steady_plus_series",
]


def solve_pde_parabolic(
    *,
    kind: ParabolicPdeKind,
    # ----- symbols -----
    x_symbol: str = "x",
    t_symbol: str = "t",
    u_symbol: str = "u",
    # ----- params -----
    k: Union[int, float, str, sp.Expr] = 1,
    L_: Union[int, float, str, sp.Expr] = sp.pi,
    n: Union[int, float, str, sp.Expr] = 1,  # for sine mode
    U0: Union[int, float, str, sp.Expr] = 1,  # for fixed ends
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Returns:
        { "problem_s": "...", "soln_s": "...", "answer_s": "..." }

    Output style:
      - HTML with \\(\\displaystyle ...\\) and aligned steps
      - answer_s is the final solution statement (mode or series form)
    """
    # -----------------------------
    # 1) Symbols + parsing
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    t = sp.Symbol(t_symbol, real=True)

    U = sp.Function(u_symbol)

    base_locals = {
        x_symbol: x,
        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,
        "tan": sp.tan,
        "asin": sp.asin,
        "acos": sp.acos,
        "atan": sp.atan,
        "sinh": sp.sinh,
        "cosh": sp.cosh,
        "tanh": sp.tanh,
        "Abs": sp.Abs,
        "sec": sp.sec,
        "csc": sp.csc,
        "cot": sp.cot,
        "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)

    K = P(k)
    L = P(L_)
    N = P(n)
    U0e = P(U0)

    if simplify:
        K = sp.simplify(K)
        L = sp.simplify(L)
        N = sp.simplify(N)
        U0e = sp.simplify(U0e)

    # -----------------------------
    # 2) Cases
    # -----------------------------
    if kind == "heat_1d_sine_mode":
        # u_t = k u_xx, u(0,t)=u(L,t)=0, u(x,0)=sin(nπx/L)
        uxt = U(x, t)
        pde_ltx = "u_t = k u_{xx}"

        bc_ltx = "u(0,t)=0,\\;u(" + Ltx(L) + ",t)=0"
        ic_ltx = "u(x,0)=\\sin\\!\\left(\\frac{" + Ltx(N) + "\\pi x}{" + Ltx(L) + "}\\right)"

        lam = (N * sp.pi / L)
        if simplify:
            lam = sp.simplify(lam)

        sol_expr = sp.sin(lam * x) * sp.exp(-K * lam**2 * t)
        if simplify:
            sol_expr = sp.simplify(sol_expr)

        answer_ltx = Ltx(uxt) + " = " + Ltx(sol_expr)

        # FIX: Replaced '0<x<' with '0 \\lt x \\lt' to avoid HTML parsing errors
        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Parabolic PDE (Heat Equation, 1D) — Single Sine Mode<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            "& \\displaystyle " + pde_ltx + ",\\; 0 \\lt x \\lt " + Ltx(L) + ",\\; t>0 \\\\"
            "& \\displaystyle \\text{BC: } " + bc_ltx + " \\\\"
            "& \\displaystyle \\text{IC: } " + ic_ltx + " \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Use separation } u=X(x)T(t) \\\\"
            "& \\displaystyle \\Rightarrow X''+\\lambda^2 X=0,\\; T'+k\\lambda^2 T=0 \\\\"
            "& \\displaystyle \\Rightarrow \\lambda=\\frac{n\\pi}{L} \\text{ from } X(0)=X(L)=0 \\\\"
            "& \\displaystyle \\Rightarrow u(x,t)=\\sin\\!\\left(\\frac{" + Ltx(N) + "\\pi x}{" + Ltx(L) + "}\\right)"
            "\\exp\\!\\left(-" + Ltx(K) + "\\left(\\frac{" + Ltx(N) + "\\pi}{" + Ltx(L) + "}\\right)^2 t\\right) \\\\"
            "\\end{aligned}\\)"
            "</div></div>"
        )

        return {
            "problem_s": ("Solve \\( " + pde_ltx + " \\) with \\( " + bc_ltx + " \\) and \\( " + ic_ltx + " \\)").strip(),
            "soln_s": soln_s.strip(),
            "answer_s": ("\\( " + answer_ltx + " \\)").strip(),
        }

    if kind == "heat_1d_fixed_ends_steady_plus_series":
        # u_t = k u_xx, u(0,t)=0, u(L,t)=U0, u(x,0)=0
        uxt = U(x, t)
        pde_ltx = "u_t = k u_{xx}"

        bc_ltx = "u(0,t)=0,\\;u(" + Ltx(L) + ",t)=" + Ltx(U0e)
        ic_ltx = "u(x,0)=0"

        # steady state: u_s = U0 * x/L
        u_s = U0e * x / L
        if simplify:
            u_s = sp.simplify(u_s)

        # v = u - u_s has homogeneous BC
        # v(x,0) = -u_s(x) = -U0 x/L
        # Fourier sine series coefficients for f(x)=x on (0,L):
        # x = sum_{m>=1} (2L*(-1)^{m+1}/(mπ)) * sin(mπx/L)
        m = sp.Symbol("m", integer=True, positive=True)
        coeff = 2 * L * (-1) ** (m + 1) / (m * sp.pi)
        v_series = sp.summation(
            (-U0e / L) * coeff * sp.sin(m * sp.pi * x / L) * sp.exp(-K * (m * sp.pi / L) ** 2 * t),
            (m, 1, sp.oo),
        )
        if simplify:
            v_series = sp.simplify(v_series)

        sol_expr = u_s + v_series

        answer_ltx = Ltx(uxt) + " = " + Ltx(sol_expr)

        # FIX: Replaced '0<x<' with '0 \\lt x \\lt' to avoid HTML parsing errors
        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Parabolic PDE (Heat Equation, 1D) — Nonzero Fixed End Temperatures<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            "& \\displaystyle " + pde_ltx + ",\\; 0 \\lt x \\lt " + Ltx(L) + " \\\\"
            "& \\displaystyle \\text{BC: } " + bc_ltx + " \\\\"
            "& \\displaystyle \\text{IC: } " + ic_ltx + " \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Steady state } u_s(x):\\; 0 = k u_s'' \\Rightarrow u_s(x)=\\frac{" + Ltx(U0e) + "}{"
            + Ltx(L) + "}x \\\\"
            "& \\displaystyle \\Rightarrow \\text{Let } v=u-u_s \\Rightarrow v_t = k v_{xx},\\; v(0,t)=v(L,t)=0 \\\\"
            "& \\displaystyle \\Rightarrow v(x,0)=-u_s(x)=-\\frac{" + Ltx(U0e) + "}{" + Ltx(L) + "}x \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Fourier sine series form:} \\\\"
            "& \\displaystyle v(x,t)=\\sum_{m=1}^{\\infty} b_m\\sin\\left(\\frac{m\\pi x}{L}\\right)e^{-k(m\\pi/L)^2 t},\\quad"
            " b_m=\\frac{2}{L}\\int_0^L v(x,0)\\sin\\left(\\frac{m\\pi x}{L}\\right)dx \\\\[10pt]"
            "& \\displaystyle \\Rightarrow u(x,t)=u_s(x)+v(x,t) \\\\"
            "\\end{aligned}\\)"
            "<div class='note'>We return the standard series solution for the transient part. In your app you can truncate the series for numeric plotting.</div>"
            "</div></div>"
        )

        return {
            "problem_s": ("Solve \\( " + pde_ltx + " \\) with \\( " + bc_ltx + " \\) and \\( " + ic_ltx + " \\)").strip(),
            "soln_s": soln_s.strip(),
            "answer_s": ("\\( " + answer_ltx + " \\)").strip(),
        }

    raise ValueError(
        "Unknown kind={!r}. Expected one of: heat_1d_sine_mode, heat_1d_fixed_ends_steady_plus_series.".format(kind)
    )