# ============================================================
# File: _solve_pde_nonhomogeneous_boundary_conditions.py
# Topic: PDE with nonhomogeneous boundary conditions
#
# Core idea: Boundary lifting / steady + transient decomposition
#
#   u(x,t) = v(x,t) + w(x,t)
#
# Choose w to satisfy the (nonhomogeneous) boundary conditions so that v has
# homogeneous BCs. Then solve v using Fourier-mode machinery (Type #8) or
# transform methods (Type #6).
#
# Implemented templates:
#
# A) Heat equation, 0<x<L
#    1) Constant Dirichlet BCs:
#         u_t = k u_xx
#         u(0,t)=A, u(L,t)=B
#         u(x,0)=f(x)
#       Choose w(x)=A + (B-A)x/L   -> v=u-w satisfies v(0,t)=v(L,t)=0
#       v_t = k v_xx, v(x,0)=f(x)-w(x)
#       v by sine series (multi-mode): Type #8
#
#    2) Time-dependent Dirichlet BCs:
#         u(0,t)=A(t), u(L,t)=B(t)
#       Choose w(x,t)=A(t)+(B(t)-A(t))x/L
#       Then v satisfies homogeneous BCs but forced PDE:
#         v_t = k v_xx - w_t
#       (We return the forced eigenfunction-expansion template; actual coefficient
#        ODEs can be built later if you want.)
#
# B) Wave equation, 0<x<L
#    1) Constant Dirichlet BCs:
#         u_tt = c^2 u_xx
#         u(0,t)=A, u(L,t)=B
#         u(x,0)=f(x), u_t(x,0)=g(x)
#       Choose w(x)=A+(B-A)x/L
#       v=u-w -> v(0,t)=v(L,t)=0
#       v_tt=c^2 v_xx
#       v(x,0)=f(x)-w(x), v_t(x,0)=g(x)
#       v by sine series (multi-mode): Type #8
#
# Notes:
#   - This file focuses on the decomposition and correct transformed IC/BC.
#   - For the "constant BC" cases we provide explicit Fourier coefficient formulas
#     (symbolic integrals) and a truncated series for readability.
#   - For time-dependent BCs we provide the canonical forced-series template
#     (leaving coefficient ODE solve as a later enhancement).
# ============================================================

from __future__ import annotations

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

import sympy as sp


PdeLiftKind = Literal[
    "heat_dirichlet_constant",
    "heat_dirichlet_time_dependent",
    "wave_dirichlet_constant",
]


def solve_pde_nonhomogeneous_boundary_conditions(
    *,
    kind: PdeLiftKind,
    # domain / parameters
    L_: Union[int, float, str, sp.Expr] = sp.pi,
    k_: Union[int, float, str, sp.Expr] = 1,
    c_: Union[int, float, str, sp.Expr] = 1,
    # boundary data
    A_bc: Union[int, float, str, sp.Expr] = 0,                 # A or A(t)
    B_bc: Union[int, float, str, sp.Expr] = 1,                 # B or B(t)
    # initial data
    f_of_x: Union[str, sp.Expr] = "x",
    g_of_x: Optional[Union[str, sp.Expr]] = None,              # for wave
    # display
    N_terms: int = 4,
    # symbols
    x_symbol: str = "x",
    t_symbol: str = "t",
    u_symbol: str = "u",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    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,
        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,
        "sinh": sp.sinh,
        "cosh": sp.cosh,
        "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)

    L = P(L_)
    k = P(k_)
    c = P(c_)
    if simplify:
        L = sp.simplify(L)
        k = sp.simplify(k)
        c = sp.simplify(c)

    A = P(A_bc)
    B = P(B_bc)
    if simplify:
        A = sp.simplify(A)
        B = sp.simplify(B)

    f = P(f_of_x) if isinstance(f_of_x, str) else sp.sympify(f_of_x, locals=base_locals)
    if simplify:
        f = sp.simplify(f)

    g = None
    if kind == "wave_dirichlet_constant":
        if g_of_x is None:
            raise ValueError("wave_dirichlet_constant requires g_of_x.")
        g = P(g_of_x) if isinstance(g_of_x, str) else sp.sympify(g_of_x, locals=base_locals)
        if simplify:
            g = sp.simplify(g)

    # Helpers for HTML
    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>"
        )

    # =========================================================
    # Constant Dirichlet lifting: w(x) = A + (B-A)x/L
    # =========================================================
    if kind in ("heat_dirichlet_constant", "wave_dirichlet_constant"):
        w = A + (B - A) * x / L
        if simplify:
            w = sp.simplify(w)

        v0 = f - w  # v(x,0)
        if simplify:
            v0 = sp.simplify(v0)

        n = sp.Symbol("n", integer=True, positive=True)

        # sine coefficients for v0
        bn = (2 / L) * sp.Integral(v0 * sp.sin(n * sp.pi * x / L), (x, 0, L))
        if simplify:
            bn = sp.simplify(bn)

        # compute first N terms (for display)
        bn_vals: List[Tuple[int, sp.Expr]] = []
        n_values = list(range(1, max(1, int(N_terms)) + 1))

        def eval_coeff(expr_in: sp.Expr, ni: int) -> sp.Expr:
            expr = expr_in.subs(n, ni)
            try:
                expr2 = expr.doit()
            except Exception:
                expr2 = expr
            return sp.simplify(expr2) if simplify else expr2

        for ni in n_values:
            bn_vals.append((ni, eval_coeff(bn, ni)))

        if kind == "heat_dirichlet_constant":
            # v(x,t)= Σ b_n e^{-k(nπ/L)^2 t} sin(...)
            v_trunc = sp.Integer(0)
            for ni, bni in bn_vals:
                v_trunc += bni * sp.exp(-k * (ni * sp.pi / L) ** 2 * t) * sp.sin(ni * sp.pi * x / L)
            if simplify:
                v_trunc = sp.simplify(v_trunc)

            u_trunc = w + v_trunc
            if simplify:
                u_trunc = sp.simplify(u_trunc)

            series_general = (
                "u(x,t)=w(x)+\\sum_{n=1}^{\\infty} b_n\\,e^{-k(\\frac{n\\pi}{L})^2 t}\\,"
                "\\sin\\left(\\frac{n\\pi x}{L}\\right)"
            )
            w_ltx = "w(x)=A+\\frac{B-A}{L}x"
            bn_ltx = "b_n=\\frac{2}{L}\\int_0^L \\big(f(x)-w(x)\\big)\\sin\\left(\\frac{n\\pi x}{L}\\right)\\,dx"

            # FIX: Replaced '0<x<' with '0 \\lt x \\lt'
            body = (
                "& \\displaystyle u_t=" + Ltx(k) + "u_{xx},\\; 0 \\lt x \\lt " + Ltx(L) + " \\\\"
                "& \\displaystyle u(0,t)=A,\\; u(L,t)=B,\\; u(x,0)=f(x) \\\\[8pt]"
                "& \\displaystyle \\Rightarrow \\text{Boundary lifting: } u=v+w \\\\"
                "& \\displaystyle \\Rightarrow " + w_ltx.replace("A", Ltx(A)).replace("B", Ltx(B)).replace("L", Ltx(L)) + " \\\\"
                "& \\displaystyle \\Rightarrow v(0,t)=v(L,t)=0 \\\\"
                "& \\displaystyle \\Rightarrow v_t=" + Ltx(k) + "v_{xx},\\; v(x,0)=f(x)-w(x) \\\\[10pt]"
                "& \\displaystyle \\Rightarrow \\text{Sine-series solution:} \\\\"
                "& \\displaystyle v(x,t)=\\sum_{n=1}^{\\infty} b_n\\,e^{-k(\\frac{n\\pi}{L})^2 t}\\sin\\left(\\frac{n\\pi x}{L}\\right) \\\\[6pt]"
                "& \\displaystyle \\Rightarrow " + bn_ltx + " \\\\[10pt]"
                "& \\displaystyle \\Rightarrow \\text{Computed first " + str(len(n_values)) + " coefficients }(b_n): \\\\"
            )
            for ni, bni in bn_vals:
                body += "& \\displaystyle b_" + str(ni) + "=" + Ltx(bni) + " \\\\"
            body += (
                "& \\\\[6pt]"
                "& \\displaystyle \\Rightarrow \\text{Truncated approximation: } u(x,t)\\approx " + Ltx(u_trunc) + " \\\\"
            )

            soln_s = card(
                "PDE Nonhomogeneous BCs — Heat (Dirichlet A,B) via Boundary Lifting",
                body,
                "We choose a simple steady boundary function \\(w\\) that matches the boundary values, so the transient \\(v\\) has homogeneous BCs and can be solved by Fourier sine series.",
            )

            problem_s = "Solve heat equation with nonzero Dirichlet BCs: \\(u_t=" + Ltx(k) + "u_{xx},\\; u(0,t)=A,\\; u(L,t)=B,\\; u(x,0)=f(x)\\)."
            answer_s = "\\( " + series_general.replace("w(x)", "w(x)") + ",\\; " + bn_ltx + ",\\; w(x)=" + Ltx(w) + " \\)"
            return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

        # -----------------------------------------------------
        # Wave constant Dirichlet BCs
        # -----------------------------------------------------
        An = bn
        Bn = (2 / (n * sp.pi * c)) * sp.Integral(g * sp.sin(n * sp.pi * x / L), (x, 0, L))
        if simplify:
            Bn = sp.simplify(Bn)

        An_vals: List[Tuple[int, sp.Expr]] = []
        Bn_vals: List[Tuple[int, sp.Expr]] = []
        for ni in n_values:
            An_vals.append((ni, eval_coeff(An, ni)))
            Bn_vals.append((ni, eval_coeff(Bn, ni)))

        v_trunc = sp.Integer(0)
        for (ni, Ani), (_, Bni) in zip(An_vals, Bn_vals):
            v_trunc += (Ani * sp.cos(ni * sp.pi * c * t / L) + Bni * sp.sin(ni * sp.pi * c * t / L)) * sp.sin(ni * sp.pi * x / L)
        if simplify:
            v_trunc = sp.simplify(v_trunc)

        u_trunc = w + v_trunc
        if simplify:
            u_trunc = sp.simplify(u_trunc)

        series_general = (
            "u(x,t)=w(x)+\\sum_{n=1}^{\\infty}\\left[A_n\\cos\\left(\\frac{n\\pi c}{L}t\\right)"
            "+B_n\\sin\\left(\\frac{n\\pi c}{L}t\\right)\\right]\\sin\\left(\\frac{n\\pi x}{L}\\right)"
        )
        w_ltx = "w(x)=A+\\frac{B-A}{L}x"
        An_ltx = "A_n=\\frac{2}{L}\\int_0^L \\big(f(x)-w(x)\\big)\\sin\\left(\\frac{n\\pi x}{L}\\right)\\,dx"
        Bn_ltx = "B_n=\\frac{2}{n\\pi c}\\int_0^L g(x)\\sin\\left(\\frac{n\\pi x}{L}\\right)\\,dx"

        # FIX: Replaced '0<x<' with '0 \\lt x \\lt'
        body = (
            "& \\displaystyle u_{tt}=" + Ltx(c**2) + "u_{xx},\\; 0 \\lt x \\lt " + Ltx(L) + " \\\\"
            "& \\displaystyle u(0,t)=A,\\; u(L,t)=B \\\\"
            "& \\displaystyle u(x,0)=f(x),\\; u_t(x,0)=g(x) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Boundary lifting: } u=v+w,\\; " + w_ltx.replace("A", Ltx(A)).replace("B", Ltx(B)).replace("L", Ltx(L)) + " \\\\"
            "& \\displaystyle \\Rightarrow v(0,t)=v(L,t)=0 \\\\"
            "& \\displaystyle \\Rightarrow v_{tt}=" + Ltx(c**2) + "v_{xx} \\\\"
            "& \\displaystyle \\Rightarrow v(x,0)=f(x)-w(x),\\; v_t(x,0)=g(x) \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Sine-series solution:} \\\\"
            "& \\displaystyle v(x,t)=\\sum_{n=1}^{\\infty}\\left[A_n\\cos\\left(\\frac{n\\pi c}{L}t\\right)+B_n\\sin\\left(\\frac{n\\pi c}{L}t\\right)\\right]\\sin\\left(\\frac{n\\pi x}{L}\\right) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow " + An_ltx + " \\\\"
            "& \\displaystyle \\Rightarrow " + Bn_ltx + " \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Computed first " + str(len(n_values)) + " coefficients:} \\\\"
        )
        for (ni, Ani), (_, Bni) in zip(An_vals, Bn_vals):
            body += "& \\displaystyle A_" + str(ni) + "=" + Ltx(Ani) + ",\\; B_" + str(ni) + "=" + Ltx(Bni) + " \\\\"
        body += (
            "& \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Truncated approximation: } u(x,t)\\approx " + Ltx(u_trunc) + " \\\\"
        )

        soln_s = card(
            "PDE Nonhomogeneous BCs — Wave (Dirichlet A,B) via Boundary Lifting",
            body,
            "For constant boundary values, pick a time-independent \\(w(x)\\) that matches the boundaries. Then \\(v=u-w\\) has homogeneous BCs and satisfies the same wave equation. Only the displacement IC shifts: \\(v(x,0)=f(x)-w(x)\\); velocity IC stays \\(g(x)\\).",
        )

        problem_s = "Solve wave equation with nonzero Dirichlet BCs: \\(u_{tt}=c^2u_{xx},\\; u(0,t)=A,\\; u(L,t)=B,\\; u(x,0)=f(x),\\; u_t(x,0)=g(x)\\)."
        answer_s = "\\( " + series_general + ",\\; " + An_ltx + ",\\; " + Bn_ltx + ",\\; w(x)=" + Ltx(w) + " \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # Time-dependent Dirichlet heat: w(x,t)=A(t)+((B(t)-A(t))/L)x
    # v_t = k v_xx - w_t  (forced with homogeneous BCs)
    # =========================================================
    if kind == "heat_dirichlet_time_dependent":
        # A and B are allowed to depend on t
        w = A + (B - A) * x / L
        if simplify:
            w = sp.simplify(w)

        wt = sp.diff(w, t)
        if simplify:
            wt = sp.simplify(wt)

        v0 = f - w.subs(t, 0)  # v(x,0)=f(x)-w(x,0)
        if simplify:
            v0 = sp.simplify(v0)

        # Provide forced-series template (do not solve coefficient ODEs fully here)
        n = sp.Symbol("n", integer=True, positive=True)
        Xn = sp.sin(n * sp.pi * x / L)
        lamn = (n * sp.pi / L) ** 2

        # Coefficient of initial condition:
        b_n0 = (2 / L) * sp.Integral(v0 * Xn, (x, 0, L))
        if simplify:
            b_n0 = sp.simplify(b_n0)

        # Forcing projection:
        # F_n(t) = (2/L) ∫_0^L (-w_t) sin(nπx/L) dx
        Fn = (2 / L) * sp.Integral((-wt) * Xn, (x, 0, L))
        if simplify:
            Fn = sp.simplify(Fn)

        # Coefficient ODE (standard): b_n'(t) + k λ_n b_n(t) = F_n(t), with b_n(0)=b_n0
        bn = sp.Function("b_n")
        # display-only symbolic:
        coeff_ode_ltx = "b_n'(t)+" + Ltx(k * lamn) + "b_n(t)=F_n(t),\\quad b_n(0)=b_{n0}"

        series_general = (
            "u(x,t)=w(x,t)+\\sum_{n=1}^{\\infty} b_n(t)\\sin\\left(\\frac{n\\pi x}{L}\\right)"
        )

        # FIX: Replaced '0<x<' with '0 \\lt x \\lt'
        body = (
            "& \\displaystyle u_t=" + Ltx(k) + "u_{xx},\\; 0 \\lt x \\lt " + Ltx(L) + " \\\\"
            "& \\displaystyle u(0,t)=A(t),\\; u(L,t)=B(t),\\; u(x,0)=f(x) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Choose } w(x,t)=A(t)+\\frac{B(t)-A(t)}{L}x \\\\"
            "& \\displaystyle \\Rightarrow v=u-w \\Rightarrow v(0,t)=v(L,t)=0 \\\\"
            "& \\displaystyle \\Rightarrow v_t = " + Ltx(k) + "v_{xx} - w_t \\\\"
            "& \\displaystyle \\Rightarrow v(x,0)=f(x)-w(x,0) \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Expand } v(x,t)=\\sum_{n=1}^{\\infty} b_n(t)\\sin\\left(\\frac{n\\pi x}{L}\\right) \\\\"
            "& \\displaystyle \\Rightarrow \\text{Mode ODE: } " + coeff_ode_ltx + " \\\\"
            "& \\displaystyle \\Rightarrow b_{n0}=\\frac{2}{L}\\int_0^L \\big(f(x)-w(x,0)\\big)\\sin\\left(\\frac{n\\pi x}{L}\\right)\\,dx \\\\"
            "& \\displaystyle \\Rightarrow F_n(t)=\\frac{2}{L}\\int_0^L \\big(-w_t(x,t)\\big)\\sin\\left(\\frac{n\\pi x}{L}\\right)\\,dx \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Final structure: } " + series_general + " \\\\"
        )

        soln_s = card(
            "PDE Nonhomogeneous BCs — Heat with Time-Dependent Dirichlet Data (Forced Modes)",
            body,
            "Time-dependent boundary lifting makes the transient equation forced by \\(-w_t\\). Each Fourier mode satisfies a first-order linear ODE in time; you can solve it by integrating factor if you want explicit \\(b_n(t)\\).",
        )

        problem_s = "Solve heat equation with time-dependent Dirichlet BCs using boundary lifting: \\(u_t=" + Ltx(k) + "u_{xx},\\; u(0,t)=A(t),\\; u(L,t)=B(t),\\; u(x,0)=f(x)\\)."
        answer_s = "\\( " + series_general + "\\;\\text{with } w(x,t)=" + Ltx(w) + ",\\; b_{n0}=" + Ltx(b_n0) + ",\\; F_n(t)=" + Ltx(Fn) + " \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError("Unknown kind={!r}. Use 'heat_dirichlet_constant', 'heat_dirichlet_time_dependent', or 'wave_dirichlet_constant'.".format(kind))