# ============================================================
# File: _solve_pde_hyperbolic.py
# Topic: Hyperbolic PDEs (Wave / Transport) — App-Friendly Canonical Cases
#
# Hyperbolic PDE hallmark (2 variables): opposite-sign principal part
#   Example: 1D wave equation  u_tt - c^2 u_xx = 0   (hyperbolic)
#
# We implement small solvers for classic hyperbolic PDE cases:
#
# (1) 1D Wave equation (d'Alembert general solution):
#       u_tt - c^2 u_xx = 0
#     General solution:
#       u(x,t) = F(x - c t) + G(x + c t)
#
# (2) 1D Advection / transport equation:
#       u_t + a u_x = 0
#     Along characteristics x - a t = const, u constant:
#       u(x,t) = F(x - a t)
#
# Notes:
#   - We return general-form solutions (F,G arbitrary).
#   - HTML output uses \\(\\displaystyle ...\\) and aligned steps.
# ============================================================

from __future__ import annotations

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

import sympy as sp


HyperbolicPdeKind = Literal[
    "wave_1d_dalembert",
    "transport_1d",
]


def solve_pde_hyperbolic(
    *,
    kind: HyperbolicPdeKind,
    # ----- symbols -----
    x_symbol: str = "x",
    t_symbol: str = "t",
    u_symbol: str = "u",
    # ----- parameters -----
    c: Union[int, float, str, sp.Expr] = 1,   # wave speed
    a: Union[int, float, str, sp.Expr] = 1,   # transport speed
    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 (general/canonical)
    """
    # -----------------------------
    # 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)

    C = P(c)
    A = P(a)

    if simplify:
        C = sp.simplify(C)
        A = sp.simplify(A)

    # -----------------------------
    # 2) Cases
    # -----------------------------
    if kind == "wave_1d_dalembert":
        # u_tt - c^2 u_xx = 0
        uxt = U(x, t)
        pde_ltx = "u_{tt} - " + Ltx(C**2) + "u_{xx} = 0"

        F = sp.Function("F")
        G = sp.Function("G")
        sol_expr = F(x - C * t) + G(x + C * t)

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

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Hyperbolic PDE (Wave Equation, 1D) — d'Alembert Form<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            "& \\displaystyle " + pde_ltx + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Characteristic variables: } \\xi=x-" + Ltx(C) + "t,\\; \\eta=x+" + Ltx(C) + "t \\\\"
            "& \\displaystyle \\Rightarrow \\text{General solution: } u(x,t)=F(\\xi)+G(\\eta) \\\\"
            "& \\displaystyle \\Rightarrow u(x,t)=F\\big(x-" + Ltx(C) + "t\\big)+G\\big(x+" + Ltx(C) + "t\\big) \\\\"
            "\\end{aligned}\\)"
            "<div class='note'>\\(F\\) and \\(G\\) are determined from initial conditions such as \\(u(x,0)\\) and \\(u_t(x,0)\\).</div>"
            "</div></div>"
        )

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

    if kind == "transport_1d":
        # u_t + a u_x = 0
        uxt = U(x, t)
        pde_ltx = "u_t + " + Ltx(A) + "u_x = 0"

        F = sp.Function("F")
        sol_expr = F(x - A * t)

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

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Hyperbolic PDE (Transport/Advection, 1D) — Characteristics<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            "& \\displaystyle " + pde_ltx + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Characteristic ODE: } \\frac{dx}{dt}=" + Ltx(A) + " \\\\"
            "& \\displaystyle \\Rightarrow x-" + Ltx(A) + "t = \\text{const} \\\\"
            "& \\displaystyle \\Rightarrow \\text{Along characteristics, } \\frac{du}{dt}=0 \\Rightarrow u=\\text{const} \\\\"
            "& \\displaystyle \\Rightarrow u(x,t)=F\\big(x-" + Ltx(A) + "t\\big) \\\\"
            "\\end{aligned}\\)"
            "<div class='note'>The profile moves rigidly with speed \\(a\\). Initial data \\(u(x,0)=f(x)\\) gives \\(u(x,t)=f(x-at)\\).</div>"
            "</div></div>"
        )

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

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