# ============================================================
# File: _solve_fourier_series_pde_solutions_multi_mode.py
# Topic: Fourier-series multi-mode solutions for PDE BVP/IVP
#
# We implement the standard textbook eigenfunction expansion engines on (0<x<L),
# focusing on full-coefficient computation (symbolic when possible; otherwise
# leave as Integral).
#
# Supported templates (homogeneous BCs):
#
# 1) Heat equation (Dirichlet–Dirichlet):
#      u_t = k u_xx,  0<x<L
#      u(0,t)=0, u(L,t)=0
#      u(x,0)=f(x)
#    Solution:
#      u(x,t)= Σ_{n=1..∞} b_n e^{-k(nπ/L)^2 t} sin(nπ x/L)
#      b_n = (2/L) ∫_0^L f(x) sin(nπ x/L) dx
#
# 2) Wave equation (Dirichlet–Dirichlet):
#      u_tt = c^2 u_xx,  0<x<L
#      u(0,t)=0, u(L,t)=0
#      u(x,0)=f(x),  u_t(x,0)=g(x)
#    Solution:
#      u(x,t)= Σ_{n=1..∞} [ A_n cos(nπ c t/L) + B_n sin(nπ c t/L) ] sin(nπ x/L)
#      A_n = (2/L) ∫_0^L f(x) sin(nπ x/L) dx
#      B_n = (2/(nπ c)) ∫_0^L g(x) sin(nπ x/L) dx
#
# 3) Laplace equation (strip y>0 with decay, Dirichlet at x=0,L):
#      u_xx + u_yy = 0, 0<x<L, y>0
#      u(0,y)=u(L,y)=0, u(x,0)=f(x), u(x,y)->0 as y->∞
#    Solution:
#      u(x,y)= Σ_{n=1..∞} b_n e^{-(nπ/L) y} sin(nπ x/L)
#      b_n = (2/L) ∫_0^L f(x) sin(nπ x/L) dx
#
# Extensions (future):
#   - Neumann–Neumann (cosine series; includes n=0)
#   - Mixed BCs (half-integer eigenvalues)
#   - Nonhomogeneous BC lifting (Type #9)
#
# Output:
#   - Returns the general series form + coefficient formulas
#   - Additionally computes first N_terms coefficients in closed-form if possible
#     and prints the truncated series as an educational approximation.
#
# Notes:
#   - SymPy may keep integrals unevaluated (still correct).
#   - We avoid heavy numerical Fourier fitting by default.
# ============================================================

from __future__ import annotations

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

import sympy as sp


FourierPdeKind = Literal[
    "heat_dd",
    "wave_dd",
    "laplace_strip_decay_dd",
]


def solve_fourier_series_pde_solutions_multi_mode(
    *,
    kind: FourierPdeKind,
    f_of_x: Union[str, sp.Expr],
    # for wave only
    g_of_x: Optional[Union[str, sp.Expr]] = None,
    # parameters
    L_: Union[int, float, str, sp.Expr] = sp.pi,
    k_: Union[int, float, str, sp.Expr] = 1,  # heat diffusivity
    c_: Union[int, float, str, sp.Expr] = 1,  # wave speed
    # truncation for display
    N_terms: int = 4,
    # symbols
    x_symbol: str = "x",
    t_symbol: str = "t",
    y_symbol: str = "y",
    u_symbol: str = "u",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Build Fourier-series multi-mode PDE solution templates with coefficient computation.

    Returns:
        { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
    """
    x = sp.Symbol(x_symbol, real=True)
    t = sp.Symbol(t_symbol, real=True)
    y = sp.Symbol(y_symbol, real=True)
    u = sp.Function(u_symbol)

    base_locals = {
        x_symbol: x,
        t_symbol: t,
        y_symbol: y,
        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)

    # Parse f(x), g(x)
    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_dd":
        if g_of_x is None:
            raise ValueError("wave_dd requires g_of_x (initial velocity).")
        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)

    # Fourier sine coefficients b_n for f(x)
    n = sp.Symbol("n", integer=True, positive=True)
    bn_formula = (2 / L) * sp.Integral(f * sp.sin(n * sp.pi * x / L), (x, 0, L))
    if simplify:
        bn_formula = sp.simplify(bn_formula)

    # For wave, also coefficient for g(x)
    An_formula = bn_formula
    Bn_formula = None
    if kind == "wave_dd":
        Bn_formula = (2 / (n * sp.pi * c)) * sp.Integral(g * sp.sin(n * sp.pi * x / L), (x, 0, L))
        if simplify:
            Bn_formula = sp.simplify(Bn_formula)

    # Compute first N_terms coefficients explicitly (n=1..N_terms)
    n_values = list(range(1, max(1, int(N_terms)) + 1))
    bn_values: List[Tuple[int, sp.Expr]] = []
    An_values: List[Tuple[int, sp.Expr]] = []
    Bn_values: List[Tuple[int, sp.Expr]] = []

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

    for ni in n_values:
        bni = eval_coeff(bn_formula, ni)
        bn_values.append((ni, bni))
        if kind == "wave_dd":
            Ani = eval_coeff(An_formula, ni)
            Bni = eval_coeff(Bn_formula, ni)
            An_values.append((ni, Ani))
            Bn_values.append((ni, Bni))

    # Truncated series for display (first N_terms)
    u_xt = u(x, t)

    if kind == "heat_dd":
        # u(x,t)=Σ b_n e^{-k(nπ/L)^2 t} sin(nπ x/L)
        terms = []
        for (ni, bni) in bn_values:
            terms.append(bni * sp.exp(-k * (ni * sp.pi / L) ** 2 * t) * sp.sin(ni * sp.pi * x / L))
        u_trunc = sp.Add(*terms) if terms else sp.Integer(0)
        if simplify:
            u_trunc = sp.simplify(u_trunc)

        series_general = (
            "u(x,t)=\\sum_{n=1}^{\\infty} b_n\\,e^{-k(\\frac{n\\pi}{L})^2 t}\\,"
            "\\sin\\left(\\frac{n\\pi x}{L}\\right)"
        )
        coeff_general = "b_n=\\frac{2}{L}\\int_0^L f(x)\\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) + ",\\; u(0,t)=u(L,t)=0,\\; u(x,0)=f(x) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Eigenfunctions: } \\sin\\left(\\frac{n\\pi x}{L}\\right) \\\\"
            "& \\displaystyle \\Rightarrow " + series_general + " \\\\[6pt]"
            "& \\displaystyle \\Rightarrow " + coeff_general + " \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Computed first " + str(len(n_values)) + " coefficients }(b_n): \\\\"
        )
        for ni, bni in bn_values:
            body += "& \\displaystyle b_" + str(ni) + "=" + Ltx(bni) + " \\\\"
        body += (
            "& \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Truncated approximation: } u(x,t)\\approx " + Ltx(u_trunc) + " \\\\"
        )

        problem_s = "Solve heat equation by Fourier series: \\(u_t=" + Ltx(k) + "u_{xx},\\; u(0,t)=u(L,t)=0,\\; u(x,0)=f(x)\\) on \\((0,L)\\)."
        answer_s = "\\( " + series_general + ",\\; " + coeff_general + " \\)"

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Fourier-Series PDE (Multi-Mode) — Heat (Dirichlet–Dirichlet)<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            + body +
            "\\end{aligned}\\)"
            "<div class='note'>If SymPy cannot evaluate the coefficient integrals, it leaves them as \\(\\int\\) (still correct). "
            "The truncated series is a practical approximation for plotting/intuition.</div>"
            "</div></div>"
        )
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    if kind == "wave_dd":
        # u(x,t)=Σ [A_n cos(nπ c t/L) + B_n sin(nπ c t/L)] sin(nπ x/L)
        terms = []
        for (ni, Ani), (_, Bni) in zip(An_values, Bn_values):
            terms.append((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))
        u_trunc = sp.Add(*terms) if terms else sp.Integer(0)
        if simplify:
            u_trunc = sp.simplify(u_trunc)

        series_general = (
            "u(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)"
        )
        An_general = "A_n=\\frac{2}{L}\\int_0^L f(x)\\sin\\left(\\frac{n\\pi x}{L}\\right)\\,dx"
        Bn_general = "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) + ",\\; u(0,t)=u(L,t)=0 \\\\[4pt]"
            "& \\displaystyle u(x,0)=f(x),\\; u_t(x,0)=g(x) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Eigenfunctions: } \\sin\\left(\\frac{n\\pi x}{L}\\right) \\\\"
            "& \\displaystyle \\Rightarrow " + series_general + " \\\\[6pt]"
            "& \\displaystyle \\Rightarrow " + An_general + " \\\\[4pt]"
            "& \\displaystyle \\Rightarrow " + Bn_general + " \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Computed first " + str(len(n_values)) + " coefficients:} \\\\"
        )
        for (ni, Ani), (_, Bni) in zip(An_values, Bn_values):
            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) + " \\\\"
        )

        problem_s = "Solve wave equation by Fourier series: \\(u_{tt}=c^2u_{xx},\\; u(0,t)=u(L,t)=0,\\; u(x,0)=f(x),\\; u_t(x,0)=g(x)\\)."
        answer_s = "\\( " + series_general + ",\\; " + An_general + ",\\; " + Bn_general + " \\)"

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Fourier-Series PDE (Multi-Mode) — Wave (Dirichlet–Dirichlet)<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            + body +
            "\\end{aligned}\\)"
            "<div class='note'>The coefficients come from orthogonality of sine eigenfunctions. "
            "If integrals remain unevaluated, they still define the correct Fourier coefficients.</div>"
            "</div></div>"
        )
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    if kind == "laplace_strip_decay_dd":
        # u_xx + u_yy = 0 on 0<x<L, y>0; u(0,y)=u(L,y)=0; u(x,0)=f(x); u->0 as y->∞
        u_xy = u(x, y)
        terms = []
        for (ni, bni) in bn_values:
            terms.append(bni * sp.exp(-(ni * sp.pi / L) * y) * sp.sin(ni * sp.pi * x / L))
        u_trunc = sp.Add(*terms) if terms else sp.Integer(0)
        if simplify:
            u_trunc = sp.simplify(u_trunc)

        series_general = (
            "u(x,y)=\\sum_{n=1}^{\\infty} b_n\\,e^{-\\frac{n\\pi}{L}y}\\,\\sin\\left(\\frac{n\\pi x}{L}\\right)"
        )
        coeff_general = "b_n=\\frac{2}{L}\\int_0^L f(x)\\sin\\left(\\frac{n\\pi x}{L}\\right)\\,dx"

        # FIX: Replaced '0<x<' with '0 \\lt x \\lt' and 'y>0' with 'y \\gt 0'
        body = (
            "& \\displaystyle u_{xx}+u_{yy}=0,\\; 0 \\lt x \\lt " + Ltx(L) + ",\\; y \\gt 0 \\\\[4pt]"
            "& \\displaystyle u(0,y)=u(L,y)=0,\\; u(x,0)=f(x),\\; u(x,y)\\to 0\\text{ as }y\\to\\infty \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Separated eigenfunctions: } \\sin\\left(\\frac{n\\pi x}{L}\\right),\\; Y_n(y)=e^{-\\frac{n\\pi}{L}y} \\\\"
            "& \\displaystyle \\Rightarrow " + series_general + " \\\\[6pt]"
            "& \\displaystyle \\Rightarrow " + coeff_general + " \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Computed first " + str(len(n_values)) + " coefficients }(b_n): \\\\"
        )
        for ni, bni in bn_values:
            body += "& \\displaystyle b_" + str(ni) + "=" + Ltx(bni) + " \\\\"
        body += (
            "& \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Truncated approximation: } u(x,y)\\approx " + Ltx(u_trunc) + " \\\\"
        )

        # FIX: Replaced '0<x<L' with '0 \\lt x \\lt L' and 'y>0' with 'y \\gt 0'
        problem_s = "Solve Laplace strip problem by Fourier series: \\(u_{xx}+u_{yy}=0\\) on \\(0 \\lt x \\lt L,\\; y \\gt 0\\) with \\(u(0,y)=u(L,y)=0\\), \\(u(x,0)=f(x)\\), \\(u\\to 0\\) as \\(y\\to\\infty\\)."
        answer_s = "\\( " + series_general + ",\\; " + coeff_general + " \\)"

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Fourier-Series PDE (Multi-Mode) — Laplace in a Strip (Decay in y)<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            + body +
            "\\end{aligned}\\)"
            "<div class='note'>This is the standard sine-series solution for Dirichlet side walls and decaying behavior as \\(y\\to\\infty\\). "
            "The coefficients are exactly the sine-series of \\(f(x)\\).</div>"
            "</div></div>"
        )
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError("Unknown kind={!r}. Use 'heat_dd', 'wave_dd', or 'laplace_strip_decay_dd'.".format(kind))