# ============================================================
# File: _solve_pde_poisson_inhomogeneous_elliptic.py
# Topic: Poisson equation & inhomogeneous elliptic PDE templates
#
# We implement a practical "textbook" template for Poisson on a rectangle:
#
#   u_xx + u_yy = f(x,y),   (x,y) in (0,L) x (0,H)
#
# with HOMOGENEOUS Dirichlet BCs:
#   u(0,y)=u(L,y)=u(x,0)=u(x,H)=0
#
# Solution by double sine eigenfunction expansion:
#
#   u(x,y) = Σ_{m=1..∞} Σ_{n=1..∞} A_{mn} sin(mπx/L) sin(nπy/H)
#
# Expand f in same basis:
#   f(x,y) = Σ Σ F_{mn} sin(mπx/L) sin(nπy/H)
#
# Coefficients:
#   F_{mn} = (4/(L H)) ∫_0^L ∫_0^H f(x,y) sin(mπx/L) sin(nπy/H) dy dx
#
# Relation:
#   A_{mn} = - F_{mn} / [ (mπ/L)^2 + (nπ/H)^2 ]
#
# Additional templates included (structural, not full solve):
#   - Nonhomogeneous boundary data: do boundary lifting u=v+w, then
#       v_xx+v_yy = f - (w_xx+w_yy) with homogeneous BCs -> same series engine
#   - Pure Neumann solvability condition (rule):
#       ∫_Ω f dΩ = ∫_{∂Ω} h ds, solution defined up to additive constant
#   - Green's function representation (statement)
#
# Output:
#   { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
#
# Notes:
#   - Uses expr.doit() (NOT sp.doit).
#   - If SymPy can’t integrate, leaves Integral(...) which is acceptable.
# ============================================================

from __future__ import annotations

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

import sympy as sp


PoissonKind = Literal[
    "poisson_rectangle_dirichlet_homogeneous",
    "poisson_rectangle_with_boundary_lifting_template",
    "poisson_neumann_solvability_rule",
    "poisson_greens_function_template",
]


def solve_pde_poisson_inhomogeneous_elliptic(
    *,
    kind: PoissonKind,
    # domain
    L_: Union[int, float, str, sp.Expr] = sp.pi,
    H_: Union[int, float, str, sp.Expr] = sp.pi,
    # source term
    f_of_xy: Union[str, sp.Expr] = "sin(x)*sin(y)",
    # truncation for display
    M_terms: int = 2,
    N_terms: int = 2,
    # symbols
    x_symbol: str = "x",
    y_symbol: str = "y",
    u_symbol: str = "u",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    x = sp.Symbol(x_symbol, real=True)
    y = sp.Symbol(y_symbol, real=True)
    u = sp.Function(u_symbol)

    base_locals = {
        x_symbol: x,
        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>"
        )

    L = P(L_)
    H = P(H_)
    if simplify:
        L = sp.simplify(L)
        H = sp.simplify(H)

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

    # =========================================================
    # 1) Main workhorse: rectangle, homogeneous Dirichlet, double sine series
    # =========================================================
    if kind == "poisson_rectangle_dirichlet_homogeneous":
        m = sp.Symbol("m", integer=True, positive=True)
        n = sp.Symbol("n", integer=True, positive=True)

        Sm = sp.sin(m * sp.pi * x / L)
        Sn = sp.sin(n * sp.pi * y / H)

        Fmn = (4 / (L * H)) * sp.Integral(sp.Integral(f * Sm * Sn, (y, 0, H)), (x, 0, L))
        if simplify:
            Fmn = sp.simplify(Fmn)

        denom = (m * sp.pi / L) ** 2 + (n * sp.pi / H) ** 2
        Amn = -Fmn / denom
        if simplify:
            Amn = sp.simplify(Amn)

        # display: compute small grid of coefficients
        Mmax = max(1, int(M_terms))
        Nmax = max(1, int(N_terms))

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

        coeffs: List[Tuple[int, int, sp.Expr, sp.Expr]] = []
        for mi in range(1, Mmax + 1):
            for ni in range(1, Nmax + 1):
                Fv = eval_coeff(Fmn, mi, ni)
                Av = eval_coeff(Amn, mi, ni)
                coeffs.append((mi, ni, Fv, Av))

        # truncated solution expression
        u_trunc = sp.Integer(0)
        for mi, ni, _, Av in coeffs:
            u_trunc += Av * sp.sin(mi * sp.pi * x / L) * sp.sin(ni * sp.pi * y / H)
        if simplify:
            u_trunc = sp.simplify(u_trunc)

        body = (
            "& \\displaystyle u_{xx}+u_{yy}=f(x,y),\\quad 0<x<" + Ltx(L) + ",\\; 0<y<" + Ltx(H) + " \\\\"
            "& \\displaystyle u(0,y)=u(" + Ltx(L) + ",y)=u(x,0)=u(x," + Ltx(H) + ")=0 \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Assume double sine series:} \\\\"
            "& \\displaystyle u(x,y)=\\sum_{m=1}^{\\infty}\\sum_{n=1}^{\\infty}A_{mn}\\sin\\left(\\frac{m\\pi x}{L}\\right)\\sin\\left(\\frac{n\\pi y}{H}\\right) \\\\[6pt]"
            "& \\displaystyle f(x,y)=\\sum_{m=1}^{\\infty}\\sum_{n=1}^{\\infty}F_{mn}\\sin\\left(\\frac{m\\pi x}{L}\\right)\\sin\\left(\\frac{n\\pi y}{H}\\right) \\\\[10pt]"
            "& \\displaystyle \\Rightarrow F_{mn}=\\frac{4}{LH}\\int_0^L\\int_0^H f(x,y)\\sin\\left(\\frac{m\\pi x}{L}\\right)\\sin\\left(\\frac{n\\pi y}{H}\\right)\\,dy\\,dx \\\\[6pt]"
            "& \\displaystyle \\Rightarrow A_{mn}=-\\frac{F_{mn}}{(\\frac{m\\pi}{L})^2+(\\frac{n\\pi}{H})^2} \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Computed (m,n) coefficients for } 1\\le m\\le " + str(Mmax) + ",\\; 1\\le n\\le " + str(Nmax) + ": \\\\"
        )
        for mi, ni, Fv, Av in coeffs:
            body += "& \\displaystyle F_{" + str(mi) + str(ni) + "}=" + Ltx(Fv) + ",\\; A_{" + str(mi) + str(ni) + "}=" + Ltx(Av) + " \\\\"
        body += (
            "& \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Truncated series approximation: } u(x,y)\\approx " + Ltx(u_trunc) + " \\\\"
        )

        soln_s = card(
            "Poisson on Rectangle — Inhomogeneous Elliptic PDE via Double Sine Expansion",
            body,
            "This is the standard solver: compute source coefficients F_{mn}, then divide by the eigenvalue (mπ/L)^2+(nπ/H)^2 to get A_{mn}.",
        )

        problem_s = "Solve Poisson equation \\(u_{xx}+u_{yy}=f(x,y)\\) on a rectangle with homogeneous Dirichlet boundary conditions."
        answer_s = (
            "\\(u(x,y)=\\sum_{m=1}^{\\infty}\\sum_{n=1}^{\\infty}A_{mn}\\sin\\left(\\frac{m\\pi x}{L}\\right)\\sin\\left(\\frac{n\\pi y}{H}\\right),\\;"
            "A_{mn}=-\\frac{F_{mn}}{(\\frac{m\\pi}{L})^2+(\\frac{n\\pi}{H})^2}\\)"
        )
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 2) Boundary lifting template for nonzero boundary data
    # =========================================================
    if kind == "poisson_rectangle_with_boundary_lifting_template":
        body = (
            "& \\displaystyle u_{xx}+u_{yy}=f(x,y) \\quad \\text{(Poisson)} \\\\"
            "& \\displaystyle \\text{Nonhomogeneous boundary data on }\\partial\\Omega: \\; u|_{\\partial\\Omega}=g \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Split: } u=u_p+u_h \\\\"
            "& \\displaystyle \\Rightarrow \\text{Choose } u_p \\text{ so that } \\nabla^2 u_p=f \\text{ (ignore boundary)} \\\\"
            "& \\displaystyle \\Rightarrow \\text{Then } u_h \\text{ satisfies Laplace: } \\nabla^2 u_h=0 \\\\"
            "& \\displaystyle \\Rightarrow \\text{Boundary for } u_h: \\; u_h|_{\\partial\\Omega}=g-u_p|_{\\partial\\Omega} \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\text{Now solve the Laplace BVP for } u_h \\text{ (your elliptic solver / Fourier series).} \\\\"
            "& \\displaystyle \\text{Alternatively (same idea): pick } w \\text{ matching boundary and set } v=u-w,\\; \\nabla^2 v=f-\\nabla^2 w. \\\\"
        )
        soln_s = card(
            "Poisson with Nonhomogeneous Boundary Data — Particular + Homogeneous Split",
            body,
            "This is the elliptic analog of boundary lifting for heat/wave. Reduce Poisson → Laplace with modified boundary data, or reduce to homogeneous BC then use double-series for the modified source.",
        )
        problem_s = "Show the standard decomposition method for Poisson with nonzero boundary data."
        answer_s = "\\(u=u_p+u_h\\) with \\(\\nabla^2 u_p=f\\) and \\(\\nabla^2 u_h=0\\), and apply BCs to \\(u_h\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 3) Neumann solvability rule
    # =========================================================
    if kind == "poisson_neumann_solvability_rule":
        body = (
            "& \\displaystyle \\nabla^2 u=f \\text{ in }\\Omega,\\quad \\frac{\\partial u}{\\partial n}=h \\text{ on }\\partial\\Omega \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Compatibility (solvability) condition:} \\\\"
            "& \\displaystyle \\int_{\\Omega} f\\,d\\Omega = \\int_{\\partial\\Omega} h\\,ds \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{If condition fails, no solution exists.} \\\\"
            "& \\displaystyle \\Rightarrow \\text{If it holds, solution is determined up to an additive constant.} \\\\"
        )
        soln_s = card(
            "Poisson with Pure Neumann BCs — Solvability Condition + Constant-Shift",
            body,
            "For pure Neumann data on a closed domain, Poisson needs a global balance condition and the solution is not unique (you fix the mean value).",
        )
        problem_s = "State the Neumann solvability condition for Poisson and the uniqueness note."
        answer_s = "\\(\\int_{\\Omega} f\\,d\\Omega = \\int_{\\partial\\Omega} h\\,ds\\), and \\(u\\) is unique up to an additive constant."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =========================================================
    # 4) Green's function representation template
    # =========================================================
    if kind == "poisson_greens_function_template":
        body = (
            "& \\displaystyle \\nabla^2 u=f \\text{ in }\\Omega \\text{ with boundary conditions} \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Green's function form:} \\\\"
            "& \\displaystyle u(\\mathbf{x})=\\int_{\\Omega} G(\\mathbf{x},\\boldsymbol{\\xi})\\,f(\\boldsymbol{\\xi})\\,d\\boldsymbol{\\xi} \\\\"
            "& \\displaystyle \\quad\\; +\\;\\text{(boundary integral terms depending on BC type)} \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{In bounded domains, }G\\text{ depends on geometry + BCs; series solutions act like discrete Green's functions.} \\\\"
        )
        soln_s = card(
            "Poisson — Green’s Function Representation (Template)",
            body,
            "This is the operator-inverse viewpoint: Green’s functions encode both the domain geometry and boundary conditions.",
        )
        problem_s = "Write the Green’s function representation for Poisson (template form)."
        answer_s = "\\(u(\\mathbf{x})=\\int_{\\Omega} 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()}

    raise ValueError(
        "Unknown kind={!r}. Use "
        "'poisson_rectangle_dirichlet_homogeneous', "
        "'poisson_rectangle_with_boundary_lifting_template', "
        "'poisson_neumann_solvability_rule', "
        "'poisson_greens_function_template'.".format(kind)
    )
