# ============================================================
# File: _solve_greens_function_methods.py
# Topic: Green’s Functions (Unified “Operator Inverse”)
#
# We implement app-friendly Green’s function templates for:
#
# A) ODE BVP (2nd order, Dirichlet–Dirichlet) in Sturm–Liouville form:
#      L[y] := (p(x) y')' + q(x) y = f(x),   x in [a,b]
#      y(a)=0, y(b)=0
#
#    Construction using two homogeneous solutions:
#      y1 solves L[y]=0 with y1(a)=0, y1'(a)≠0
#      y2 solves L[y]=0 with y2(b)=0, y2'(b)≠0
#
#    Wronskian-like:
#      W(ξ) = p(ξ) ( y1(ξ) y2'(ξ) - y1'(ξ) y2(ξ) )
#
#    Green’s function:
#      G(x,ξ) = { y1(x) y2(ξ)/W(ξ) , a<=x<=ξ
#               { y1(ξ) y2(x)/W(ξ) , ξ<=x<=b
#
#    Solution:
#      y(x) = ∫_a^b G(x,ξ) f(ξ) dξ
#
# B) Classic ODE example:
#      y'' = f(x), y(0)=0, y(L)=0
#    => y1=x, y2=L-x, W=-L
#    => G(x,ξ)= x(L-ξ)/L for x<=ξ,  ξ(L-x)/L for x>=ξ
#
# C) PDE kernels (templates):
#    - Poisson free-space (2D log / 3D 1/r) statement
#    - Heat kernel on R (whole line)
#    - Wave (retarded) kernel statement
#
# Notes:
#   - For general p(x),q(x) we do NOT attempt to solve y1,y2 automatically;
#     we present the symbolic construction and let the caller provide y1,y2.
#   - For the classic y'' example we fully build G(x,ξ).
# ============================================================

from __future__ import annotations

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

import sympy as sp


GreensKind = Literal[
    "ode_bvp_dirichlet_template",
    "ode_bvp_ypp_classic_kernel",
    "pde_poisson_fundamental_solution_template",
    "pde_heat_kernel_template",
    "pde_wave_retarded_template",
]


def solve_greens_function_methods(
    *,
    kind: GreensKind,
    # symbols
    x_symbol: str = "x",
    xi_symbol: str = "xi",
    t_symbol: str = "t",
    tau_symbol: str = "tau",
    y_symbol: str = "y",
    u_symbol: str = "u",
    # interval params for ODE
    a_: Union[int, float, str, sp.Expr] = 0,
    b_: Union[int, float, str, sp.Expr] = 1,
    # for classic y'' example
    L_: Optional[Union[int, float, str, sp.Expr]] = None,
    # operator coefficients for template
    p_of_x: Union[int, float, str, sp.Expr] = 1,
    q_of_x: Union[int, float, str, sp.Expr] = 0,
    f_of_x: Union[str, sp.Expr] = "f(x)",
    # optional provided homogeneous solutions for template
    y1_of_x: Optional[Union[str, sp.Expr]] = None,
    y2_of_x: Optional[Union[str, sp.Expr]] = None,
    # PDE params (heat/wave)
    k_: Union[int, float, str, sp.Expr] = 1,
    c_: Union[int, float, str, sp.Expr] = 1,
    # locals
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    x = sp.Symbol(x_symbol, real=True)
    xi = sp.Symbol(xi_symbol, real=True)
    t = sp.Symbol(t_symbol, real=True)
    tau = sp.Symbol(tau_symbol, real=True)

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

    base_locals = {
        x_symbol: x,
        xi_symbol: xi,
        t_symbol: t,
        tau_symbol: tau,
        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)

    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>"
        )

    a = P(a_)
    b = P(b_)
    p = P(p_of_x)
    q = P(q_of_x)
    k = P(k_)
    c = P(c_)

    if simplify:
        a = sp.simplify(a)
        b = sp.simplify(b)
        p = sp.simplify(p)
        q = sp.simplify(q)
        k = sp.simplify(k)
        c = sp.simplify(c)

    # =========================================================
    # A) General ODE BVP Dirichlet–Dirichlet template
    # =========================================================
    if kind == "ode_bvp_dirichlet_template":
        fexpr = P(f_of_x) if isinstance(f_of_x, str) else sp.sympify(f_of_x, locals=base_locals)
        if simplify:
            fexpr = sp.simplify(fexpr)

        y1 = None
        y2 = None
        if y1_of_x is not None:
            y1 = P(y1_of_x) if isinstance(y1_of_x, str) else sp.sympify(y1_of_x, locals=base_locals)
        if y2_of_x is not None:
            y2 = P(y2_of_x) if isinstance(y2_of_x, str) else sp.sympify(y2_of_x, locals=base_locals)

        # leave y1,y2 symbolic if not provided
        if y1 is None:
            y1 = sp.Function("y_1")(x)
        if y2 is None:
            y2 = sp.Function("y_2")(x)

        # W(ξ) = p(ξ)( y1(ξ) y2'(ξ) - y1'(ξ) y2(ξ) )
        y1_xi = y1.subs(x, xi)
        y2_xi = y2.subs(x, xi)
        y1p_xi = sp.diff(y1, x).subs(x, xi)
        y2p_xi = sp.diff(y2, x).subs(x, xi)

        Wxi = p.subs(x, xi) * (y1_xi * y2p_xi - y1p_xi * y2_xi)
        if simplify:
            Wxi = sp.simplify(Wxi)

        G_piece = sp.Piecewise(
            (y1 * y2_xi / Wxi, sp.Le(x, xi)),
            (y1_xi * y2 / Wxi, True),
        )

        # solution integral
        y_sol = sp.Integral(G_piece * fexpr.subs(x, xi), (xi, a, b))

        body = (
            "& \\displaystyle L[y]=(p(x)y')'+q(x)y=f(x),\\quad " + Ltx(a) + "\\le x\\le " + Ltx(b) + " \\\\"
            "& \\displaystyle y(" + Ltx(a) + ")=0,\\; y(" + Ltx(b) + ")=0 \\qquad (\\text{Dirichlet–Dirichlet}) \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Pick two homogeneous solutions:} \\\\"
            "& \\displaystyle L[y_1]=0,\\; y_1(" + Ltx(a) + ")=0; \\qquad L[y_2]=0,\\; y_2(" + Ltx(b) + ")=0 \\\\[8pt]"
            "& \\displaystyle \\Rightarrow W(\\xi)=p(\\xi)\\big(y_1(\\xi)y_2'(\\xi)-y_1'(\\xi)y_2(\\xi)\\big)=" + Ltx(Wxi) + " \\\\[10pt]"
            "& \\displaystyle \\Rightarrow G(x,\\xi)=\\begin{cases}"
            "\\dfrac{y_1(x)\\,y_2(\\xi)}{W(\\xi)}, & " + Ltx(a) + "\\le x\\le \\xi \\\\"
            "\\dfrac{y_1(\\xi)\\,y_2(x)}{W(\\xi)}, & \\xi\\le x\\le " + Ltx(b) +
            "\\end{cases} \\\\[12pt]"
            "& \\displaystyle \\Rightarrow y(x)=\\int_{" + Ltx(a) + "}^{" + Ltx(b) + "} G(x,\\xi)\\,f(\\xi)\\,d\\xi \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Compact (Piecewise) }G(x,\\xi):\\; " + Ltx(G_piece) + " \\\\"
            "& \\displaystyle \\Rightarrow \\text{Solution: } y(x)= " + Ltx(y_sol) + " \\\\"
        )

        soln_s = card(
            "Green’s Function — ODE BVP (Dirichlet–Dirichlet) Construction Template",
            body,
            "G is continuous at x=ξ but has a controlled derivative jump so that L[G]=δ(x-ξ). Provide y1,y2 explicitly to get a closed-form kernel.",
        )
        problem_s = "Construct the Green’s function for a 2nd-order ODE BVP and write the solution integral."
        answer_s = "\\(y(x)=\\int_a^b G(x,\\xi)f(\\xi)d\\xi\\) with \\(G\\) built from two homogeneous solutions (piecewise)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # B) Classic ODE kernel: y''=f(x), y(0)=0, y(L)=0
    # =========================================================
    if kind == "ode_bvp_ypp_classic_kernel":
        # interval: [0,L] by default unless user sets a,b
        if L_ is not None:
            L = P(L_)
            if simplify:
                L = sp.simplify(L)
            a0 = sp.Integer(0)
            b0 = L
        else:
            a0 = a
            b0 = b
            L = b0 - a0

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

        # For [0,L]: y1=x, y2=L-x. For general [a,b] shift: y1=x-a, y2=b-x.
        y1 = x - a0
        y2 = b0 - x
        W = sp.simplify(y1 * sp.diff(y2, x) - sp.diff(y1, x) * y2) if simplify else (y1 * sp.diff(y2, x) - sp.diff(y1, x) * y2)
        # For 0..L: W = -L; for a..b: W=-(b-a)
        Wc = sp.simplify(W)

        G = sp.Piecewise(
            (y1 * (b0 - xi) / (-Wc), sp.Le(x, xi)),   # equals (x-a)(b-ξ)/(b-a)
            ((xi - a0) * y2 / (-Wc), True),           # equals (ξ-a)(b-x)/(b-a)
        )
        if simplify:
            G = sp.simplify(G)

        y_sol = sp.Integral(G * fexpr.subs(x, xi), (xi, a0, b0))

        body = (
            "& \\displaystyle y''=f(x),\\quad y(" + Ltx(a0) + ")=0,\\; y(" + Ltx(b0) + ")=0 \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Choose } y_1(x)=x-" + Ltx(a0) + ",\\; y_2(x)=" + Ltx(b0) + "-x \\\\[6pt]"
            "& \\displaystyle \\Rightarrow W=y_1y_2'-y_1'y_2=" + Ltx(Wc) + " \\\\[10pt]"
            "& \\displaystyle \\Rightarrow G(x,\\xi)=\\begin{cases}"
            "\\dfrac{(x-" + Ltx(a0) + ")(" + Ltx(b0) + "-\\xi)}{" + Ltx(b0 - a0) + "}, & x\\le \\xi \\\\"
            "\\dfrac{(\\xi-" + Ltx(a0) + ")(" + Ltx(b0) + "-x)}{" + Ltx(b0 - a0) + "}, & x\\ge \\xi"
            "\\end{cases} \\\\[10pt]"
            "& \\displaystyle \\Rightarrow y(x)=\\int_{" + Ltx(a0) + "}^{" + Ltx(b0) + "} G(x,\\xi)\\,f(\\xi)\\,d\\xi \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Piecewise }G(x,\\xi):\\; " + Ltx(G) + " \\\\"
            "& \\displaystyle \\Rightarrow \\text{Solution: } y(x)= " + Ltx(y_sol) + " \\\\"
        )

        soln_s = card(
            "Green’s Function — Classic ODE BVP Kernel (y''=f, Dirichlet ends)",
            body,
            "This kernel is a reusable impulse-response operator inverse for the second derivative on a finite interval with zero endpoint values.",
        )
        problem_s = "Solve \\(y''=f(x)\\) with Dirichlet BCs using a Green’s function on a finite interval."
        answer_s = "\\(y(x)=\\int_a^b G(x,\\xi)f(\\xi)d\\xi\\) with \\(G\\) as the piecewise linear kernel."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # C) PDE Poisson fundamental solution template
    # =========================================================
    if kind == "pde_poisson_fundamental_solution_template":
        r = sp.Symbol("r", positive=True, real=True)
        body = (
            "& \\displaystyle \\nabla^2 u=f \\quad (\\text{Poisson}) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Free-space fundamental solution (kernel):} \\\\"
            "& \\displaystyle \\text{In 3D: } G(\\mathbf{x},\\boldsymbol{\\xi})\\propto \\frac{1}{\\lVert\\mathbf{x}-\\boldsymbol{\\xi}\\rVert} \\\\"
            "& \\displaystyle \\text{In 2D: } G(\\mathbf{x},\\boldsymbol{\\xi})\\propto \\ln\\lVert\\mathbf{x}-\\boldsymbol{\\xi}\\rVert \\\\[8pt]"
            "& \\displaystyle \\Rightarrow u(\\mathbf{x})=\\int_{\\Omega} G(\\mathbf{x},\\boldsymbol{\\xi})\\,f(\\boldsymbol{\\xi})\\,d\\boldsymbol{\\xi} + \\text{boundary terms} \\\\"
            "& \\displaystyle \\Rightarrow \\text{In bounded domains, }G\\text{ must also enforce the boundary conditions (often via eigenfunction expansion).} \\\\"
        )
        soln_s = card(
            "Green’s Functions — Poisson Fundamental Solution (Template)",
            body,
            "On bounded domains, the Green’s function is modified to satisfy the boundary conditions; series solutions often act as the discrete Green’s function.",
        )
        problem_s = "State the Green’s-function/fundamental-solution viewpoint for Poisson."
        answer_s = "\\(u(\\mathbf{x})=\\int G(\\mathbf{x},\\boldsymbol{\\xi}) f(\\boldsymbol{\\xi}) d\\boldsymbol{\\xi}+\\text{boundary terms}.\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # D) PDE heat kernel (whole line) template
    # =========================================================
    if kind == "pde_heat_kernel_template":
        # G(x,t;xi,tau) = 1/sqrt(4πk(t-τ)) exp(-(x-ξ)^2/(4k(t-τ))) for t>τ
        G = (1 / sp.sqrt(4 * sp.pi * k * (t - tau))) * sp.exp(-((x - xi) ** 2) / (4 * k * (t - tau)))
        body = (
            "& \\displaystyle u_t - " + Ltx(k) + "u_{xx}=f(x,t),\\quad x\\in\\mathbb{R} \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Heat kernel (impulse response):} \\\\"
            "& \\displaystyle G(x,t;\\xi,\\tau)=\\frac{1}{\\sqrt{4\\pi k(t-\\tau)}}\\exp\\left(-\\frac{(x-\\xi)^2}{4k(t-\\tau)}\\right),\\; t>\\tau \\\\[8pt]"
            "& \\displaystyle \\Rightarrow u(x,t)=\\int_{-\\infty}^{\\infty}G(x,t;\\xi,0)\\,\\phi(\\xi)\\,d\\xi"
            "+\\int_0^t\\int_{-\\infty}^{\\infty}G(x,t;\\xi,\\tau)\\,f(\\xi,\\tau)\\,d\\xi\\,d\\tau \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Kernel expression: } " + Ltx(G) + " \\\\"
        )
        soln_s = card(
            "Green’s Functions — Heat Kernel on \\(\\mathbb{R}\\) (Template)",
            body,
            "This is the convolution/impulse-response engine for diffusion: initial data and forcing are propagated by the same kernel.",
        )
        problem_s = "Write the heat-equation Green’s function on the whole line and the solution representation."
        answer_s = "\\(u(x,t)=\\int G(x,t;\\xi,0)\\phi(\\xi)d\\xi + \\int_0^t\\int G(x,t;\\xi,\\tau)f(\\xi,\\tau)d\\xi d\\tau\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # E) PDE wave (retarded kernel) template (statement)
    # =========================================================
    if kind == "pde_wave_retarded_template":
        body = (
            "& \\displaystyle u_{tt}-" + Ltx(c**2) + "u_{xx}=f(x,t) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\text{Wave Green’s function encodes finite propagation speed (retarded kernel).} \\\\"
            "& \\displaystyle \\Rightarrow \\text{Solution is a spacetime convolution with }G\\text{: } u=G * f \\text{ plus terms from initial data.} \\\\"
            "& \\displaystyle \\Rightarrow \\text{(In 1D, the homogeneous part reduces to d'Alembert form on }\\mathbb{R}\\text{).} \\\\"
        )
        soln_s = card(
            "Green’s Functions — Wave Equation (Retarded Kernel) Template",
            body,
            "For the wave equation, the Green’s function has support only inside the light cone (signals propagate at speed c).",
        )
        problem_s = "State the Green’s-function viewpoint for the forced wave equation (template)."
        answer_s = "\\(u(x,t)=\\iint G(x,t;\\xi,\\tau)f(\\xi,\\tau)d\\xi d\\tau + \\text{(initial data terms)}\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError(
        "Unknown kind={!r}. Use "
        "'ode_bvp_dirichlet_template', 'ode_bvp_ypp_classic_kernel', "
        "'pde_poisson_fundamental_solution_template', 'pde_heat_kernel_template', "
        "'pde_wave_retarded_template'.".format(kind)
    )
