# ============================================================
# File: _solve_transform_methods_fourier_laplace.py
# Topic: Transform Methods (Fourier/Laplace) for ODE & PDE
#
# This solver is template-driven (app-friendly) and focuses on the reusable
# transform patterns:
#
# A) ODE (LTI) as impulse-response / convolution engine:
#      y'' + a y' + b y = f(t)
#    Transfer function (zero IC): H(s)=1/(s^2 + a s + b)
#    Impulse response: h(t)=L^{-1}{H(s)}
#    Output: y(t) = (h * f)(t) = ∫_0^t h(t-τ) f(τ) dτ
#    Special inputs:
#      - impulse at t=c: f(t)=δ(t-c) -> y(t)=h(t-c)H(t-c)
#      - step at t=c:    f(t)=H(t-c) -> y(t)=∫_0^{t-c} h(u) du · H(t-c)
#
# B) PDE (Heat) on whole line via Fourier transform -> heat kernel convolution:
#      u_t = k u_xx, x∈R, u(x,0)=φ(x)
#      u(x,t)=∫_{-∞}^{∞} φ(ξ) (1/√(4πkt)) exp(-(x-ξ)^2/(4kt)) dξ
#
# C) PDE (Wave) on whole line via Fourier transform -> d'Alembert form:
#      u_tt = c^2 u_xx, u(x,0)=φ(x), u_t(x,0)=ψ(x)
#      u(x,t)=1/2[φ(x-ct)+φ(x+ct)] + (1/(2c))∫_{x-ct}^{x+ct} ψ(ξ)dξ
#
# D) PDE (Heat) on half-line x>0 using sine/cosine transform equivalently via
#    odd/even extension (closed kernels):
#    Dirichlet u(0,t)=0:
#      u(x,t)=∫_0^∞ φ(ξ) (1/√(4πkt)) [e^{-(x-ξ)^2/(4kt)} - e^{-(x+ξ)^2/(4kt)}] dξ
#    Neumann u_x(0,t)=0:
#      u(x,t)=∫_0^∞ φ(ξ) (1/√(4πkt)) [e^{-(x-ξ)^2/(4kt)} + e^{-(x+ξ)^2/(4kt)}] dξ
#
# Notes:
#   - We avoid fragile symbolic Fourier/Laplace evaluation; return canonical
#     convolution/kernel forms (exact and standard).
#   - HTML output uses \\(\\displaystyle ...\\) aligned steps.
# ============================================================

from __future__ import annotations

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

import sympy as sp


TransformKind = Literal[
    "ode_lti_convolution",
    "ode_lti_impulse_shift",
    "ode_lti_step_shift",
    "heat_whole_line_kernel",
    "wave_whole_line_dalembert",
    "heat_half_line_dirichlet",
    "heat_half_line_neumann",
]


def solve_transform_methods(
    *,
    kind: TransformKind,
    # symbols
    t_symbol: str = "t",
    x_symbol: str = "x",
    xi_symbol: str = "xi",
    tau_symbol: str = "tau",
    y_symbol: str = "y",
    u_symbol: str = "u",
    # ODE params
    a: Union[int, float, str, sp.Expr] = 0,
    b: Union[int, float, str, sp.Expr] = 1,
    shift_c: Union[int, float, str, sp.Expr] = 1,  # step/impulse time shift
    # PDE params
    k: Union[int, float, str, sp.Expr] = 1,        # heat diffusivity
    c: Union[int, float, str, sp.Expr] = 1,        # wave speed
    # locals
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Returns:
        { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
    """
    t = sp.Symbol(t_symbol, real=True)
    x = sp.Symbol(x_symbol, real=True)
    xi = sp.Symbol(xi_symbol, real=True)
    tau = sp.Symbol(tau_symbol, real=True)

    y = sp.Function(y_symbol)
    u = sp.Function(u_symbol)

    # Data functions for PDE
    phi = sp.Function("phi")
    psi = sp.Function("psi")
    f = sp.Function("f")  # generic forcing for ODE

    base_locals = {
        t_symbol: t,
        x_symbol: x,
        xi_symbol: xi,
        tau_symbol: tau,
        y_symbol: y,
        u_symbol: u,
        "phi": phi,
        "psi": psi,
        "f": f,
        "pi": sp.pi,
        "E": sp.E,
        "exp": sp.exp,
        "log": sp.log,
        "ln": sp.log,
        "sqrt": sp.sqrt,
        "sin": sp.sin,
        "cos": sp.cos,
        "Abs": sp.Abs,
        "Heaviside": sp.Heaviside,
        "DiracDelta": sp.DiracDelta,
        "Integral": sp.Integral,
    }
    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)

    A = P(a)
    B = P(b)
    K = P(k)
    C = P(c)
    C0 = P(shift_c)

    if simplify:
        A = sp.simplify(A)
        B = sp.simplify(B)
        K = sp.simplify(K)
        C = sp.simplify(C)
        C0 = sp.simplify(C0)

    # ------------- helpers
    def card(title: str, body_aligned_ltx: 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_ltx +
            "\\end{aligned}\\)"
            + note_html +
            "</div></div>"
        )

    # =========================================================
    # A) ODE LTI convolution / impulse / step
    # =========================================================
    if kind in ("ode_lti_convolution", "ode_lti_impulse_shift", "ode_lti_step_shift"):
        y_t = y(t)

        # ODE: y'' + a y' + b y = input(t)
        # Transfer function for zero IC: H(s)=1/(s^2 + a s + b)
        s = sp.Symbol("s", real=True, positive=True)
        Hs = 1 / (s**2 + A*s + B)
        if simplify:
            Hs = sp.simplify(Hs)

        # Define h(t) formally as inverse Laplace of H(s)
        # (We keep it symbolic to avoid version issues)
        h = sp.Function("h")
        h_def = sp.Eq(h(t), sp.Function("Linv")(Hs, t))  # display-only placeholder

        # Use convolution form
        conv = sp.Integral(h(t - tau) * f(tau), (tau, 0, t))

        ode_ltx = "y'' + " + Ltx(A) + "y' + " + Ltx(B) + "y = "
        input_ltx = ""
        sol_ltx = ""
        note = (
            "Laplace turns an LTI ODE into algebra: \\(Y(s)=H(s)F(s)\\) (for zero initial conditions). "
            "Then \\(y(t)=(h*f)(t)\\)."
        )

        if kind == "ode_lti_convolution":
            input_ltx = "f(t)"
            sol_ltx = Ltx(y_t) + " = \\int_0^t h(t-\\tau) f(\\tau)\\,d\\tau"
            body = (
                "& \\displaystyle " + ode_ltx + input_ltx + " \\\\[6pt]"
                "& \\displaystyle \\Rightarrow \\text{Transfer function: } H(s)=\\frac{1}{s^2+" + Ltx(A) + "s+" + Ltx(B) + "} \\\\[6pt]"
                "& \\displaystyle \\Rightarrow \\text{Impulse response: } h(t)=\\mathcal{L}^{-1}\\{H(s)\\} \\\\[6pt]"
                "& \\displaystyle \\Rightarrow \\text{Convolution solution (zero IC): } " + sol_ltx + " \\\\"
            )
            problem_s = "Solve using transforms: \\( y''+" + Ltx(A) + "y'+" + Ltx(B) + "y=f(t) \\) (use convolution form)."
            answer_s = "\\( " + sol_ltx + " \\)"

        elif kind == "ode_lti_impulse_shift":
            # input = δ(t-c)
            input_ltx = "\\delta(t-" + Ltx(C0) + ")"
            # output: h(t-c)H(t-c)
            sol_ltx = Ltx(y_t) + " = h\\big(t-" + Ltx(C0) + "\\big)\\,\\mathrm{H}\\big(t-" + Ltx(C0) + "\\big)"
            body = (
                "& \\displaystyle " + ode_ltx + input_ltx + " \\\\[6pt]"
                "& \\displaystyle \\Rightarrow H(s)=\\frac{1}{s^2+" + Ltx(A) + "s+" + Ltx(B) + "} \\\\[6pt]"
                "& \\displaystyle \\Rightarrow \\text{If } f(t)=\\delta(t-c),\\; \\text{then } y(t)=h(t-c)\\,\\mathrm{H}(t-c) \\\\[6pt]"
                "& \\displaystyle \\Rightarrow " + sol_ltx + " \\\\"
            )
            problem_s = "Solve using Laplace/impulse response: \\( y''+" + Ltx(A) + "y'+" + Ltx(B) + "y=\\delta(t-" + Ltx(C0) + ") \\)."
            answer_s = "\\( " + sol_ltx + " \\)"
            note = note + " Here \\(\\mathrm{H}\\) denotes the Heaviside step; the impulse response shifted gives the output."

        else:
            # step input = H(t-c)
            input_ltx = "\\mathrm{H}(t-" + Ltx(C0) + ")"
            # y(t)= ∫_0^{t-c} h(u)du * H(t-c)  (use change u=t-τ)
            uvar = sp.Symbol("u", real=True)
            step_out = sp.Integral(h(uvar), (uvar, 0, t - C0)) * sp.Heaviside(t - C0)
            sol_ltx = Ltx(y_t) + " = \\left(\\int_0^{t-c} h(u)\\,du\\right)\\mathrm{H}(t-c)"
            body = (
                "& \\displaystyle " + ode_ltx + input_ltx + " \\\\[6pt]"
                "& \\displaystyle \\Rightarrow H(s)=\\frac{1}{s^2+" + Ltx(A) + "s+" + Ltx(B) + "} \\\\[6pt]"
                "& \\displaystyle \\Rightarrow \\text{Step response is the integral of impulse response:} \\\\[4pt]"
                "& \\displaystyle \\Rightarrow y(t)=\\int_0^t h(t-\\tau)\\,\\mathrm{H}(\\tau-c)\\,d\\tau \\\\[4pt]"
                "& \\displaystyle \\Rightarrow " + sol_ltx + " \\\\"
            )
            problem_s = "Solve using Laplace/step response: \\( y''+" + Ltx(A) + "y'+" + Ltx(B) + "y=\\mathrm{H}(t-" + Ltx(C0) + ") \\)."
            answer_s = "\\( " + sol_ltx + " \\)"
            note = note + " Step response = convolution of \\(h\\) with Heaviside; equivalently integrate \\(h\\) after the shift."

        soln_s = card("Transform Methods — LTI ODE as Convolution/Impulse Response", body, note)
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # B) Heat equation on whole line (kernel)
    # =========================================================
    if kind == "heat_whole_line_kernel":
        u_xt = u(x, t)
        kernel = (1 / sp.sqrt(4 * sp.pi * K * t)) * sp.exp(-(x - xi) ** 2 / (4 * K * t))
        if simplify:
            kernel = sp.simplify(kernel)

        sol_expr = sp.Integral(phi(xi) * kernel, (xi, -sp.oo, sp.oo))

        body = (
            "& \\displaystyle u_t = " + Ltx(K) + "u_{xx},\\; x\\in\\mathbb{R},\\; u(x,0)=\\phi(x) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Fourier transform in }x:\\; \\widehat{u}(\\omega,t)=\\mathcal{F}\\{u\\} \\\\"
            "& \\displaystyle \\Rightarrow \\widehat{u}_t = -" + Ltx(K) + "\\omega^2\\widehat{u} \\\\"
            "& \\displaystyle \\Rightarrow \\widehat{u}(\\omega,t)=\\widehat{\\phi}(\\omega)e^{-" + Ltx(K) + "\\omega^2 t} \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Inverse transform gives heat-kernel convolution:} \\\\[4pt]"
            "& \\displaystyle " + Ltx(u_xt) + " = " + Ltx(sol_expr) + " \\\\"
        )

        note = "This is the standard whole-line heat kernel solution: \\(u=\\phi*G\\) with Gaussian kernel."
        soln_s = card("Transform Methods — Heat Equation on \\(\\mathbb{R}\\) (Fourier Kernel)", body, note)

        problem_s = "Solve using Fourier transform: \\(u_t=" + Ltx(K) + "u_{xx},\\; x\\in\\mathbb{R},\\; u(x,0)=\\phi(x)\\)."
        answer_s = "\\( " + Ltx(u_xt) + " = " + Ltx(sol_expr) + " \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # C) Wave equation on whole line (d'Alembert)
    # =========================================================
    if kind == "wave_whole_line_dalembert":
        u_xt = u(x, t)
        left = x - C * t
        right = x + C * t
        sol_expr = sp.Rational(1, 2) * (phi(left) + phi(right)) + (sp.Rational(1, 2) / C) * sp.Integral(psi(xi), (xi, left, right))
        if simplify:
            sol_expr = sp.simplify(sol_expr)

        body = (
            "& \\displaystyle u_{tt} = " + Ltx(C**2) + "u_{xx},\\; x\\in\\mathbb{R} \\\\[4pt]"
            "& \\displaystyle u(x,0)=\\phi(x),\\; u_t(x,0)=\\psi(x) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Fourier transform in }x\\text{ gives } \\widehat{u}_{tt} + " + Ltx(C**2) + "\\omega^2\\widehat{u}=0 \\\\"
            "& \\displaystyle \\Rightarrow \\text{Invert to obtain d'Alembert form:} \\\\[4pt]"
            "& \\displaystyle " + Ltx(u_xt) + " = " + Ltx(sol_expr) + " \\\\"
        )

        note = "This whole-line wave solution can be derived via Fourier transform (or characteristics)."
        soln_s = card("Transform Methods — Wave Equation on \\(\\mathbb{R}\\) (Fourier → d'Alembert)", body, note)

        problem_s = "Solve (whole-line) wave equation: \\(u_{tt}=" + Ltx(C**2) + "u_{xx},\\; u(x,0)=\\phi(x),\\; u_t(x,0)=\\psi(x)\\)."
        answer_s = "\\( " + Ltx(u_xt) + " = " + Ltx(sol_expr) + " \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # D) Heat equation on half-line via sine/cosine (kernel difference/sum)
    # =========================================================
    if kind in ("heat_half_line_dirichlet", "heat_half_line_neumann"):
        u_xt = u(x, t)
        base = 1 / sp.sqrt(4 * sp.pi * K * t)
        g1 = sp.exp(-(x - xi) ** 2 / (4 * K * t))
        g2 = sp.exp(-(x + xi) ** 2 / (4 * K * t))
        if simplify:
            base = sp.simplify(base)
            g1 = sp.simplify(g1)
            g2 = sp.simplify(g2)

        if kind == "heat_half_line_dirichlet":
            sol_expr = sp.Integral(phi(xi) * base * (g1 - g2), (xi, 0, sp.oo))
            bc = "u(0,t)=0"
            title = "Transform Methods — Heat on Half-Line (Dirichlet, Sine / Odd Extension)"
            note = "Dirichlet at 0 corresponds to a sine transform; equivalently use odd extension giving kernel difference."
        else:
            sol_expr = sp.Integral(phi(xi) * base * (g1 + g2), (xi, 0, sp.oo))
            bc = "u_x(0,t)=0"
            title = "Transform Methods — Heat on Half-Line (Neumann, Cosine / Even Extension)"
            note = "Neumann at 0 corresponds to a cosine transform; equivalently use even extension giving kernel sum."

        if simplify:
            sol_expr = sp.simplify(sol_expr)

        body = (
            "& \\displaystyle u_t = " + Ltx(K) + "u_{xx},\\; x>0,\\; " + bc + ",\\; u(x,0)=\\phi(x) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Use sine/cosine transform (or odd/even extension)} \\\\[6pt]"
            "& \\displaystyle \\Rightarrow " + Ltx(u_xt) + " = " + Ltx(sol_expr) + " \\\\"
        )

        soln_s = card(title, body, note)
        problem_s = "Solve half-line heat equation: \\(u_t=" + Ltx(K) + "u_{xx},\\; x>0,\\; " + bc + ",\\; u(x,0)=\\phi(x)\\)."
        answer_s = "\\( " + Ltx(u_xt) + " = " + Ltx(sol_expr) + " \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError(
        "Unknown kind={!r}. Expected one of: ode_lti_convolution, ode_lti_impulse_shift, ode_lti_step_shift, "
        "heat_whole_line_kernel, wave_whole_line_dalembert, heat_half_line_dirichlet, heat_half_line_neumann.".format(kind)
    )
