# ============================================================
# File: _solve_integral_equations.py
# Topic: Integral Equations (Volterra / Fredholm; 1st/2nd kind)
#
# Covers solver-templates (symbolic + algorithmic):
#   1) Volterra / Fredholm, first kind vs second kind
#   2) Successive approximations (Picard / Neumann series iteration)
#   3) Degenerate (separable) kernels -> finite-dimensional linear system
#   4) Volterra convolution kernel -> Laplace transform solution template
#   5) Numerical quadrature discretization -> linear system
#
# This is "app-friendly": provides canonical forms + reduction steps.
# It can compute some symbolic pieces (integrals/linear systems) when possible,
# but it is not meant to be a full CAS integral-equation solver.
#
# Output:
#   { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
#
# IMPORTANT SYMPY NOTE:
#   SymPy's parser uses Python syntax; the bare identifier "lambda" is a Python
#   keyword and cannot be parsed via sympify("lambda"). We therefore map string
#   inputs "lambda"/"Λ" to a SymPy Symbol explicitly.
# ============================================================

from __future__ import annotations

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

import sympy as sp


IntegralEqKind = Literal[
    "volterra_second_kind_neumann_iteration",
    "fredholm_second_kind_neumann_iteration",
    "fredholm_second_kind_degenerate_kernel",
    "volterra_convolution_laplace_template",
    "fredholm_quadrature_discretization_template",
]


def solve_integral_equations(
    *,
    kind: IntegralEqKind,
    # Common symbols / intervals
    x_symbol: str = "x",
    t_symbol: str = "t",
    a_: Union[int, float, str, sp.Expr] = 0,
    b_: Union[int, float, str, sp.Expr] = 1,
    lam_: Union[int, float, str, sp.Expr] = "lambda",
    # Function strings for display (kept generic if not supplied)
    f_str: str = "f(x)",
    kernel_str: str = "K(x,t)",
    # Degenerate kernel specification: K(x,t)=sum a_k(x)b_k(t)
    # Provide lists of expressions/strings for a_k(x), b_k(t)
    degenerate_a_list: Optional[List[Union[str, sp.Expr]]] = None,
    degenerate_b_list: Optional[List[Union[str, sp.Expr]]] = None,
    # Convolution kernel for Volterra: y(t)=f(t)+∫_0^t k(t-τ) y(τ) dτ
    conv_k_str: str = "k(t)",
    # Iteration depth for Neumann/Picard display
    n_iter: int = 3,
    # Discretization size for quadrature template
    N_grid: int = 5,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    # -----------------------------
    # 1) Symbols + locals
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    t = sp.Symbol(t_symbol, real=True)
    s = sp.Symbol("s", real=True)

    # Reserve a safe symbol object for lambda-parameter
    lam_sym = sp.Symbol("lambda", real=True)

    base_locals: Dict[str, Any] = {
        x_symbol: x,
        t_symbol: t,
        "s": s,
        "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,
        "Gamma": sp.gamma,
        "Heaviside": sp.Heaviside,
        "DiracDelta": sp.DiracDelta,
        # IMPORTANT: map 'lambda' to a SymPy Symbol, but never sympify("lambda")
        "lambda": lam_sym,
        "lam": lam_sym,
        "Λ": lam_sym,
    }
    if locals_map:
        base_locals.update(locals_map)

    def P(v: Union[int, float, str, sp.Expr]) -> sp.Expr:
        """
        Safe parser:
          - SymPy cannot parse the bare string 'lambda' because it's a Python keyword.
          - So if a user passes 'lambda' (or variants), map to our Symbol directly.
        """
        if isinstance(v, str):
            vv = v.strip()
            if vv in {"lambda", "Λ"}:
                return lam_sym
            if vv == "":
                return sp.Integer(0)
            return sp.sympify(vv, locals=base_locals)
        return sp.sympify(v, locals=base_locals)

    # Parse interval endpoints
    a = P(a_)
    b = P(b_)

    # Parse lambda parameter safely
    if isinstance(lam_, (int, float)):
        lam_v = sp.sympify(lam_)
    else:
        lam_v = P(lam_)
    if simplify:
        lam_v = sp.simplify(lam_v)

    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) Neumann / successive approximations for Volterra 2nd kind
    # y(x)=f(x)+λ∫_a^x K(x,t) y(t) dt
    # ------------------------------------------------------------
    if kind == "volterra_second_kind_neumann_iteration":
        # Use displayed lambda symbol if user supplied something else
        lam_disp = Ltx(lam_v) if lam_v is not None else Ltx(lam_sym)

        body = (
            "& \\displaystyle \\text{Volterra (2nd kind): } y(x)= " + f_str +
            " + " + lam_disp + "\\int_{a}^{x} " + kernel_str + "\\,y(t)\\,dt \\\\[10pt]"
            "& \\displaystyle \\textbf{Successive approximations (Picard/Neumann):} \\\\[2pt]"
            "& \\displaystyle y_0(x)=f(x) \\\\"
        )
        for k in range(1, max(1, int(n_iter)) + 1):
            body += (
                "& \\displaystyle y_" + str(k) + "(x)=f(x)+"
                + lam_disp + "\\int_a^x " + kernel_str + "\\,y_" + str(k - 1) + "(t)\\,dt \\\\"
            )
        body += (
            "\\\\[-2pt]"
            "& \\displaystyle \\Rightarrow \\text{If it converges: } y = f + \\lambda K[f] + \\lambda^2 K^2[f] + \\cdots \\\\"
        )

        soln_s = card(
            "Integral Equations — Volterra (2nd kind) via Neumann / Picard Iteration",
            body,
            "Volterra equations are “causal”: upper limit is x. Discretized systems are typically triangular (fast to solve). Convergence depends on kernel size / |λ| (operator norm).",
        )
        problem_s = "Solve a Volterra integral equation of the second kind using successive approximations."
        answer_s = "\\(y_{n+1}(x)=f(x)+\\lambda\\int_a^x K(x,t)y_n(t)dt\\), starting from \\(y_0=f\\) (Neumann series if convergent)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # B) Neumann / successive approximations for Fredholm 2nd kind
    # y(x)=f(x)+λ∫_a^b K(x,t) y(t) dt
    # ------------------------------------------------------------
    if kind == "fredholm_second_kind_neumann_iteration":
        lam_disp = Ltx(lam_v) if lam_v is not None else Ltx(lam_sym)

        body = (
            "& \\displaystyle \\text{Fredholm (2nd kind): } y(x)= " + f_str +
            " + " + lam_disp + "\\int_{a}^{b} " + kernel_str + "\\,y(t)\\,dt \\\\[10pt]"
            "& \\displaystyle \\textbf{Successive approximations (Neumann):} \\\\[2pt]"
            "& \\displaystyle y_0(x)=f(x) \\\\"
        )
        for k in range(1, max(1, int(n_iter)) + 1):
            body += (
                "& \\displaystyle y_" + str(k) + "(x)=f(x)+"
                + lam_disp + "\\int_a^b " + kernel_str + "\\,y_" + str(k - 1) + "(t)\\,dt \\\\"
            )
        body += (
            "\\\\[-2pt]"
            "& \\displaystyle \\Rightarrow \\text{Formal series: } y = f + \\lambda K[f] + \\lambda^2 K^2[f] + \\cdots \\\\"
        )

        soln_s = card(
            "Integral Equations — Fredholm (2nd kind) via Neumann Series",
            body,
            "Fredholm equations use fixed limits. Quadrature discretization yields a dense linear system. Neumann series converges when |λ| is small enough in an operator sense.",
        )
        problem_s = "Solve a Fredholm integral equation of the second kind using successive approximations."
        answer_s = "\\(y_{n+1}(x)=f(x)+\\lambda\\int_a^b K(x,t)y_n(t)dt\\), starting from \\(y_0=f\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # C) Degenerate (separable) kernel reduction (Fredholm 2nd kind)
    # K(x,t)= Σ_{k=1..N} a_k(x) b_k(t)
    # y(x)=f(x)+λ∫_a^b K(x,t)y(t)dt
    # -> y(x)=f(x)+λ Σ a_k(x) c_k,  c_k=∫ b_k(t) y(t) dt
    # -> linear system for c
    # ------------------------------------------------------------
    if kind == "fredholm_second_kind_degenerate_kernel":
        lam_disp = Ltx(lam_v) if lam_v is not None else Ltx(lam_sym)

        if not degenerate_a_list or not degenerate_b_list or (len(degenerate_a_list) != len(degenerate_b_list)):
            # Provide a safe default N=2 example
            degenerate_a_list = ["a1(x)", "a2(x)"]
            degenerate_b_list = ["b1(t)", "b2(t)"]

        N = len(degenerate_a_list)

        # Ensure parsing is safe (but these are mainly for display)
        _ = [P(expr) for expr in degenerate_a_list]
        _ = [P(expr) for expr in degenerate_b_list]

        kernel_sum = " + ".join(
            ["(" + str(degenerate_a_list[i]) + ")(" + str(degenerate_b_list[i]) + ")" for i in range(N)]
        )

        body = (
            "& \\displaystyle \\text{Fredholm (2nd kind): } y(x)= " + f_str +
            " + " + lam_disp + "\\int_{a}^{b} K(x,t)\\,y(t)\\,dt \\\\[8pt]"
            "& \\displaystyle \\text{Assume degenerate kernel: } K(x,t)=\\sum_{k=1}^{N} a_k(x)b_k(t) \\\\"
            "& \\displaystyle \\Rightarrow K(x,t)= " + kernel_sum + " \\quad (N=" + str(N) + ") \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\int_a^b K(x,t)y(t)dt = \\sum_{k=1}^{N} a_k(x)\\,c_k \\\\"
            "& \\displaystyle \\text{where } c_k=\\int_a^b b_k(t)\\,y(t)\\,dt \\\\[10pt]"
            "& \\displaystyle \\Rightarrow y(x)=f(x)+" + lam_disp + "\\sum_{k=1}^{N} a_k(x)c_k \\\\[10pt]"
            "& \\displaystyle \\text{Now multiply both sides by }b_j(x)\\text{ and integrate }x\\in[a,b]: \\\\"
        )

        body += (
            "& \\displaystyle c_j = \\int_a^b b_j(x)f(x)\\,dx + "
            + lam_disp + "\\sum_{k=1}^{N} c_k\\int_a^b b_j(x)a_k(x)\\,dx \\quad (j=1..N) \\\\"
        )
        body += "& \\displaystyle \\Rightarrow \\text{This is an }N\\times N\\text{ linear system for }(c_1,\\dots,c_N). \\\\"
        body += "& \\displaystyle \\Rightarrow \\text{Then } y(x)=f(x)+" + lam_disp + "\\sum_{k=1}^{N} a_k(x)c_k. \\\\"

        soln_s = card(
            "Integral Equations — Degenerate Kernel (Separable) Reduction",
            body,
            "Degenerate kernels turn the integral operator into finite rank. After introducing constants c_k, you solve an N×N linear system, then reconstruct y(x).",
        )
        problem_s = "Solve a Fredholm 2nd-kind integral equation with a degenerate (separable) kernel."
        answer_s = "\\(K(x,t)=\\sum a_k(x)b_k(t)\\Rightarrow y(x)=f(x)+\\lambda\\sum a_k(x)c_k\\) with \\(c\\) from an \\(N\\times N\\) linear system."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # D) Volterra convolution kernel -> Laplace transform template
    # y(t)=f(t)+∫_0^t k(t-τ)y(τ)dτ
    # -> Y(s)=F(s)+K(s)Y(s) => Y(s)=F(s)/(1-K(s))
    # ------------------------------------------------------------
    if kind == "volterra_convolution_laplace_template":
        body = (
            "& \\displaystyle \\text{Volterra convolution form: } y(t)=f(t)+\\int_0^t "
            + conv_k_str + "(t-\\tau)\\,y(\\tau)\\,d\\tau \\\\[10pt]"
            "& \\displaystyle \\textbf{Laplace transform: } \\mathcal{L}\\{y\\}=Y(s),\\ \\mathcal{L}\\{f\\}=F(s),\\ \\mathcal{L}\\{k\\}=K(s) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\mathcal{L}\\left\\{\\int_0^t k(t-\\tau)y(\\tau)d\\tau\\right\\}=K(s)Y(s) \\\\[8pt]"
            "& \\displaystyle \\Rightarrow Y(s)=F(s)+K(s)Y(s) \\\\[4pt]"
            "& \\displaystyle \\Rightarrow Y(s)=\\frac{F(s)}{1-K(s)} \\\\[8pt]"
            "& \\displaystyle \\Rightarrow y(t)=\\mathcal{L}^{-1}\\left\\{\\frac{F(s)}{1-K(s)}\\right\\}. \\\\"
        )
        soln_s = card(
            "Integral Equations — Volterra Convolution via Laplace Transform",
            body,
            "This is the cleanest solvable class: convolution becomes multiplication in Laplace space. It mirrors LTI impulse-response / transfer-function logic.",
        )
        problem_s = "Solve a Volterra convolution integral equation using Laplace transforms."
        answer_s = "\\(Y(s)=\\frac{F(s)}{1-K(s)}\\Rightarrow y(t)=\\mathcal{L}^{-1}\\{F(s)/(1-K(s))\\}.\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # E) Numerical quadrature discretization template (Fredholm)
    # Discretize integral to linear system.
    # ------------------------------------------------------------
    if kind == "fredholm_quadrature_discretization_template":
        lam_disp = Ltx(lam_v) if lam_v is not None else Ltx(lam_sym)

        N = max(3, int(N_grid))
        body = (
            "& \\displaystyle \\text{Fredholm (2nd kind): } y(x)=f(x)+"
            + lam_disp + "\\int_a^b " + kernel_str + "\\,y(t)\\,dt \\\\[10pt]"
            "& \\displaystyle \\textbf{Quadrature discretization:} \\\\[2pt]"
            "& \\displaystyle \\text{Choose nodes } x_i\\ (i=1..N)\\text{ on }[a,b]\\text{ and weights }w_j \\\\"
            "& \\displaystyle \\int_a^b K(x_i,t)y(t)dt \\approx \\sum_{j=1}^{N} w_j\\,K(x_i,x_j)\\,y(x_j) \\\\[10pt]"
            "& \\displaystyle \\Rightarrow y_i = f_i + " + lam_disp + "\\sum_{j=1}^{N} w_j\\,K_{ij}\\,y_j \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\sum_{j=1}^{N}\\big(\\delta_{ij}-" + lam_disp + " w_j K_{ij}\\big)y_j=f_i \\\\[10pt]"
            "& \\displaystyle \\text{This is a linear system } A\\mathbf{y}=\\mathbf{f}\\text{ (dense in general).} \\\\"
        )
        soln_s = card(
            "Integral Equations — Quadrature Discretization (Fredholm)",
            body,
            "Volterra discretizations are usually triangular (faster). Fredholm discretizations are typically dense, but still manageable for moderate N.",
        )
        problem_s = "Show how to discretize a Fredholm 2nd-kind integral equation using quadrature."
        answer_s = "\\(\\sum_j(\\delta_{ij}-\\lambda w_jK(x_i,x_j))y_j=f(x_i)\\) (a linear system)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError(
        "Unknown kind={!r}. Use one of: "
        "'volterra_second_kind_neumann_iteration',"
        "'fredholm_second_kind_neumann_iteration',"
        "'fredholm_second_kind_degenerate_kernel',"
        "'volterra_convolution_laplace_template',"
        "'fredholm_quadrature_discretization_template'.".format(kind)
    )
