# ============================================================
# File: _solve_pde_separation_of_variables_1d_dirichlet.py
# Topic: PDE Separation of Variables (1D Heat/Wave, Dirichlet BCs)
#
# Heat:
#   u_t = k u_xx,   0 < x < L, t > 0
#   u(0,t)=0, u(L,t)=0
#   u(x,0)=f(x)
#
# Wave:
#   u_tt = c^2 u_xx, 0 < x < L, t > 0
#   u(0,t)=0, u(L,t)=0
#   u(x,0)=f(x),  u_t(x,0)=g(x)
#
# With Dirichlet BCs:
#   X_n(x) = sin(n*pi*x/L), lambda_n = (n*pi/L)^2
#
# This solver is optimized for "single-mode" ICs:
#   f(x)=A sin(n*pi*x/L) and for wave optionally g(x)=B sin(n*pi*x/L),
# so the final answer is a single term (no Fourier coefficient integrals).
# ============================================================

from __future__ import annotations

from typing import Any, Dict, Optional, Union

import sympy as sp


def solve_pde_separation_of_variables_1d_dirichlet(
    *,
    pde_kind: str,
    L_value: Union[str, int, float, sp.Expr] = 1,
    k_value: Optional[Union[str, int, float, sp.Expr]] = None,
    c_value: Optional[Union[str, int, float, sp.Expr]] = None,
    mode_n: int = 1,
    A_value: Union[str, int, float, sp.Expr] = 1,
    B_value: Optional[Union[str, int, float, sp.Expr]] = None,
    x_symbol: str = "x",
    t_symbol: str = "t",
    u_symbol: str = "u",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Solve 1D Heat/Wave PDE via separation of variables for a single eigenmode.

    pde_kind: "heat" or "wave"

    Heat IC:
      u(x,0)=A sin(n*pi*x/L)

    Wave IC:
      u(x,0)=A sin(n*pi*x/L)
      u_t(x,0)=B sin(n*pi*x/L)  (optional; if omitted, assume B=0)

    Returns:
      { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
    """

    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,
        "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,
        "Piecewise": sp.Piecewise,
        "Heaviside": sp.Heaviside,
    }
    if locals_map:
        base_locals.update(locals_map)

    def _parse(v):
        if isinstance(v, str):
            return sp.sympify(v, locals=base_locals)
        return sp.sympify(v)

    def L(expr: sp.Expr) -> str:
        return sp.latex(expr)

    pde_kind = (pde_kind or "").strip().lower()
    if pde_kind not in ("heat", "wave"):
        raise ValueError("pde_kind must be 'heat' or 'wave'.")

    Lsym = _parse(L_value)
    Lsym = sp.simplify(Lsym) if simplify else Lsym

    n = int(mode_n)
    if n <= 0:
        raise ValueError("mode_n must be a positive integer.")

    A = _parse(A_value)
    A = sp.simplify(A) if simplify else A

    if pde_kind == "heat":
        if k_value is None:
            k_value = 1
        k = _parse(k_value)
        k = sp.simplify(k) if simplify else k
    else:
        if c_value is None:
            c_value = 1
        c = _parse(c_value)
        c = sp.simplify(c) if simplify else c
        if B_value is None:
            B_value = 0
        B = _parse(B_value)
        B = sp.simplify(B) if simplify else B

    # Eigenvalue + eigenfunction
    lam_n = (n * sp.pi / Lsym) ** 2
    lam_n = sp.simplify(lam_n) if simplify else lam_n

    Xn = sp.sin(n * sp.pi * x / Lsym)
    Xn = sp.simplify(Xn) if simplify else Xn

    # Nice printable u(x,t) symbols (avoid u(x,t)_t style)
    u_xt = "u(x,t)"
    x_ltx = L(x)
    t_ltx = L(t)
    L_ltx = L(Lsym)

    # -----------------------------
    # HEAT
    # -----------------------------
    if pde_kind == "heat":

        # Heat: T' + k*lambda*T = 0 => T = exp(-k*lambda*t)
        Tn = sp.exp(-k * lam_n * t)
        Tn = sp.simplify(Tn) if simplify else Tn
        u_ans = sp.simplify(A * Xn * Tn) if simplify else (A * Xn * Tn)

        pde_ltx = f"\\frac{{\\partial {u_xt}}}{{\\partial {t_ltx}}} = {L(k)}\\,\\frac{{\\partial^2 {u_xt}}}{{\\partial {x_ltx}^2}}"
        bc_ltx = f"{u_xt}\\big|_{{{x_ltx}=0}}=0,\\; {u_xt}\\big|_{{{x_ltx}={L_ltx}}}=0"
        ic_ltx = f"{u_xt}\\big|_{{{t_ltx}=0}}={L(A)}\\sin\\left(\\frac{{{n}\\pi {x_ltx}}}{{{L_ltx}}}\\right)"

        soln_core = f"""
         Note: PDE Separation of Variables — 1D Heat (Dirichlet)

            \\(\\displaystyle
            \\begin{{aligned}}
            & \\displaystyle {pde_ltx},\\quad 0<{x_ltx}<{L_ltx},\\; {t_ltx}>0 \\\\
            & \\displaystyle {bc_ltx},\\quad {ic_ltx} \\\\[10pt]

            & \\displaystyle \\Rightarrow \\text{{Assume }} u(x,t)=X(x)T(t). \\\\
            & \\displaystyle \\Rightarrow u_t = XT',\\; u_{{xx}}=X''T \\\\
            & \\displaystyle \\Rightarrow XT' = kX''T \\\\
            & \\displaystyle \\Rightarrow \\frac{{T'}}{{kT}} = \\frac{{X''}}{{X}} = -\\lambda \\\\[10pt]

            & \\displaystyle \\Rightarrow X''+\\lambda X=0,\\quad X(0)=0,\\; X(L)=0 \\\\
            & \\displaystyle \\Rightarrow \\lambda_n=\\left(\\frac{{n\\pi}}{{L}}\\right)^2,\\quad X_n(x)=\\sin\\left(\\frac{{n\\pi x}}{{L}}\\right) \\\\[8pt]

            & \\displaystyle \\Rightarrow T' + k\\lambda_n T = 0 \\\\
            & \\displaystyle \\Rightarrow T_n(t)=\\exp\\left(-k\\lambda_n t\\right) \\\\[10pt]

            & \\displaystyle \\Rightarrow u(x,t)=A\\,X_n(x)T_n(t) = {L(u_ans)} \\\\
            \\end{{aligned}}
            \\)

        """.replace("\n", "")

        problem_s = f"Solve the IVP by separation of variables: \\( {pde_ltx} \\), with \\( {bc_ltx} \\) and \\( {ic_ltx} \\)."
        answer_s = f"\\( {u_xt} = {L(u_ans)} \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_core.strip(), "answer_s": answer_s.strip()}

    # -----------------------------
    # WAVE
    # -----------------------------
    # omega_n = c * sqrt(lambda_n) = c * (n*pi/L)  (clean, always real for L>0)
    omega_n = sp.simplify(c * n * sp.pi / Lsym) if simplify else (c * n * sp.pi / Lsym)

    # T = A cos(omega t) + (B/omega) sin(omega t)
    Tn = A * sp.cos(omega_n * t) + (B / omega_n) * sp.sin(omega_n * t)
    Tn = sp.simplify(Tn) if simplify else Tn

    u_ans = sp.simplify(Xn * Tn) if simplify else (Xn * Tn)

    pde_ltx = f"\\frac{{\\partial^2 {u_xt}}}{{\\partial {t_ltx}^2}} = {L(c)}^2\\,\\frac{{\\partial^2 {u_xt}}}{{\\partial {x_ltx}^2}}"
    bc_ltx = f"{u_xt}\\big|_{{{x_ltx}=0}}=0,\\; {u_xt}\\big|_{{{x_ltx}={L_ltx}}}=0"
    ic1_ltx = f"{u_xt}\\big|_{{{t_ltx}=0}}={L(A)}\\sin\\left(\\frac{{{n}\\pi {x_ltx}}}{{{L_ltx}}}\\right)"

    # IMPORTANT FIX: if B=0, print just u_t(x,0)=0 (not "0 sin(...)")
    if simplify and sp.simplify(B) == 0:
        ic2_ltx = f"\\frac{{\\partial {u_xt}}}{{\\partial {t_ltx}}}\\Big|_{{{t_ltx}=0}}=0"
        Tn_ltx = "T_n(t)=A\\cos(\\omega_n t)"
    else:
        ic2_ltx = f"\\frac{{\\partial {u_xt}}}{{\\partial {t_ltx}}}\\Big|_{{{t_ltx}=0}}={L(B)}\\sin\\left(\\frac{{{n}\\pi {x_ltx}}}{{{L_ltx}}}\\right)"
        Tn_ltx = "T_n(t)=A\\cos(\\omega_n t) + \\frac{B}{\\omega_n}\\sin(\\omega_n t)"

    soln_core = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: PDE Separation of Variables — 1D Wave (Dirichlet)<br/></div>

        \\(\\displaystyle
        \\begin{{aligned}}
        & \\displaystyle {pde_ltx},\\quad 0<{x_ltx}<{L_ltx},\\; {t_ltx}>0 \\\\
        & \\displaystyle {bc_ltx} \\\\
        & \\displaystyle {ic1_ltx},\\quad {ic2_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Assume }} u(x,t)=X(x)T(t). \\\\
        & \\displaystyle \\Rightarrow u_{{tt}}=XT'',\\; u_{{xx}}=X''T \\\\
        & \\displaystyle \\Rightarrow XT'' = c^2X''T \\\\
        & \\displaystyle \\Rightarrow \\frac{{T''}}{{c^2T}} = \\frac{{X''}}{{X}} = -\\lambda \\\\[10pt]

        & \\displaystyle \\Rightarrow X''+\\lambda X=0,\\quad X(0)=0,\\; X(L)=0 \\\\
        & \\displaystyle \\Rightarrow \\lambda_n=\\left(\\frac{{n\\pi}}{{L}}\\right)^2,\\quad X_n(x)=\\sin\\left(\\frac{{n\\pi x}}{{L}}\\right) \\\\[8pt]

        & \\displaystyle \\Rightarrow T'' + c^2\\lambda_n T = 0 \\\\
        & \\displaystyle \\Rightarrow \\omega_n = c\\sqrt{{\\lambda_n}} = {L(omega_n)} \\\\
        & \\displaystyle \\Rightarrow {Tn_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow u(x,t)=X_n(x)T_n(t) = {L(u_ans)} \\\\
        \\end{{aligned}}
        \\)
      </div>
    </div>
    """.replace("\n", "")

    problem_s = f"Solve the IVP by separation of variables: \\( {pde_ltx} \\), with \\( {bc_ltx} \\), \\( {ic1_ltx} \\), and \\( {ic2_ltx} \\)."
    answer_s = f"\\( {u_xt} = {L(u_ans)} \\)"

    return {"problem_s": problem_s.strip(), "soln_s": soln_core.strip(), "answer_s": answer_s.strip()}
