# ============================================================
# File: _solve_pde_linear.py
# Topic: Linear PDE (Common Canonical Forms)
#
# We implement a small, app-friendly solver for a few standard linear PDEs:
#
# 1) First-order linear homogeneous PDE with constant coefficients:
#       a u_x + b u_y = 0
#    Method of characteristics:
#       dy/dx = b/a  =>  a y - b x = const
#       u is constant along characteristics
#    General solution:
#       u(x,y) = F(a y - b x)
#
# 2) 1D Wave equation (linear hyperbolic):
#       u_tt - c^2 u_xx = 0
#    d'Alembert solution:
#       u(x,t) = F(x - c t) + G(x + c t)
#
# 3) 1D Heat equation (linear parabolic), single sine-mode IBVP:
#       u_t = k u_xx,  0<x<L
#       u(0,t)=u(L,t)=0
#       u(x,0)=sin(nπx/L)
#    Solution:
#       u(x,t)=sin(nπx/L) * exp(-k (nπ/L)^2 t)
#
# Notes:
#   - We intentionally 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


LinearPdeKind = Literal["first_order_const_coeff", "wave_1d", "heat_1d_sine_mode"]


def solve_pde_linear(
    *,
    kind: LinearPdeKind,
    # ----- symbols -----
    x_symbol: str = "x",
    y_symbol: str = "y",
    t_symbol: str = "t",
    u_symbol: str = "u",
    # ----- kind-specific params -----
    a: Union[int, float, str, sp.Expr] = 1,   # for first order: a u_x + b u_y = 0
    b: Union[int, float, str, sp.Expr] = 1,
    c: Union[int, float, str, sp.Expr] = 1,   # for wave
    k: Union[int, float, str, sp.Expr] = 1,   # for heat
    n: Union[int, float, str, sp.Expr] = 1,   # for heat mode
    L_: Union[int, float, str, sp.Expr] = sp.pi,  # for heat mode
    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-form)
    """
    # -----------------------------
    # 1) Symbols + parsing
    # -----------------------------
    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,
        "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)

    # Parse parameters
    A = P(a)
    B = P(b)
    Cc = P(c)
    Kk = P(k)
    Nn = P(n)
    LL = P(L_)

    if simplify:
        A = sp.simplify(A)
        B = sp.simplify(B)
        Cc = sp.simplify(Cc)
        Kk = sp.simplify(Kk)
        Nn = sp.simplify(Nn)
        LL = sp.simplify(LL)

    # -----------------------------
    # 2) Solve by kind
    # -----------------------------
    if kind == "first_order_const_coeff":
        # PDE: a u_x + b u_y = 0
        uxy = U(x, y)
        pde_ltx = f"{Ltx(A)}\\,u_x + {Ltx(B)}\\,u_y = 0"

        # characteristic invariant: A*y - B*x = const (for A != 0)
        # General solution: u(x,y)=F(A*y - B*x)
        s = sp.Symbol("s", real=True)
        F = sp.Function("F")

        if A == 0:
            # then B u_y = 0 -> u_y=0 -> u = F(x)
            invariant = x
            sol_expr = F(x)
            char_step = f"{Ltx(B)}\\,u_y=0 \\Rightarrow u_y=0"
            inv_step = f"\\Rightarrow u(x,y)=F({Ltx(x)})"
        else:
            invariant = sp.simplify(A * y - B * x) if simplify else (A * y - B * x)
            # dy/dx = B/A
            char_step = f"\\frac{{dy}}{{dx}}=\\frac{{{Ltx(B)}}}{{{Ltx(A)}}}"
            inv_step = f"{Ltx(A)}y - {Ltx(B)}x = \\text{{const}}"
            sol_expr = F(invariant)

        answer_ltx = f"{Ltx(uxy)} = {Ltx(sol_expr)}"

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Linear PDE (First-Order, Constant Coefficients)<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            f"& \\displaystyle {pde_ltx} \\\\[8pt]"
            f"& \\displaystyle \\Rightarrow \\text{{Characteristics: }} {char_step} \\\\"
            f"& \\displaystyle \\Rightarrow \\text{{Invariant: }} {inv_step} \\\\[6pt]"
            f"& \\displaystyle \\Rightarrow \\text{{General solution: }} u(x,y)=F\\big({Ltx(invariant)}\\big) \\\\"
            "\\end{aligned}\\)"
            "<div class='note'>Here \\(F\\) is an arbitrary differentiable function (set by boundary/initial data).</div>"
            "</div></div>"
        )

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

    if kind == "wave_1d":
        # PDE: u_tt - c^2 u_xx = 0
        uxt = U(x, t)
        pde_ltx = f"u_{{tt}} - {Ltx(Cc**2)}\\,u_{{xx}} = 0"

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

        answer_ltx = f"{Ltx(uxt)} = {Ltx(sol_expr)}"

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

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

    if kind == "heat_1d_sine_mode":
        # IBVP single-mode:
        # u_t = k u_xx, 0<x<L, u(0,t)=u(L,t)=0, u(x,0)=sin(nπx/L)
        uxt = U(x, t)
        lam = sp.simplify((Nn * sp.pi / LL) ** 2) if simplify else (Nn * sp.pi / LL) ** 2

        pde_ltx = f"u_t = {Ltx(Kk)}\\,u_{{xx}}"
        ic_ltx = f"u(x,0)=\\sin\\!\\left(\\frac{{{Ltx(Nn)}\\pi x}}{{{Ltx(LL)}}}\\right)"
        bc_ltx = f"u(0,t)=0,\\;u({Ltx(LL)},t)=0"

        sol_expr = sp.sin(Nn * sp.pi * x / LL) * sp.exp(-Kk * lam * t)
        answer_ltx = f"{Ltx(uxt)} = {Ltx(sol_expr)}"

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Linear PDE (Heat Equation, 1D) — Single Sine Mode<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            f"& \\displaystyle {pde_ltx},\\; 0<x<{Ltx(LL)} \\\\"
            f"& \\displaystyle \\text{{BC: }} {bc_ltx} \\\\"
            f"& \\displaystyle \\text{{IC: }} {ic_ltx} \\\\[10pt]"
            f"& \\displaystyle \\Rightarrow \\text{{Try separated mode }} u(x,t)=X(x)T(t) \\\\"
            f"& \\displaystyle \\Rightarrow X''+\\lambda X=0,\\; T'+{Ltx(Kk)}\\lambda T=0 \\\\"
            f"& \\displaystyle \\Rightarrow \\lambda=\\left(\\frac{{{Ltx(Nn)}\\pi}}{{{Ltx(LL)}}}\\right)^2 \\\\"
            f"& \\displaystyle \\Rightarrow u(x,t)=\\sin\\!\\left(\\frac{{{Ltx(Nn)}\\pi x}}{{{Ltx(LL)}}}\\right)"
            f"\\exp\\!\\left(-{Ltx(Kk)}\\left(\\frac{{{Ltx(Nn)}\\pi}}{{{Ltx(LL)}}}\\right)^2 t\\right) \\\\"
            "\\end{aligned}\\)"
            "</div></div>"
        )

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

    raise ValueError(f"Unknown kind={kind!r}. Expected one of: first_order_const_coeff, wave_1d, heat_1d_sine_mode.")
