# ============================================================
# File: _solve_integro_differential_equations.py
# Topic: Integro-Differential Equations (Volterra/Fredholm)
#
# Covers standard templates:
#   1) Volterra integro-differential -> convert to integral equation + iteration
#   2) Linear convolution case -> Laplace transform template
#   3) Degenerate kernel -> convert to augmented ODE system template
#   4) Numerical "memory" stepping (history sum) template
#   5) Fredholm integro-differential template (fixed limits)
#
# Output:
#   { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
#
# Notes:
#   - We do not attempt to fully "solve" arbitrary integro-differential equations
#     in closed form. Instead we return canonical transformations and templates.
#   - IMPORTANT: SymPy cannot parse the bare string "lambda" (Python keyword).
#     We map "lambda"/"Λ" to a SymPy Symbol explicitly.
# ============================================================

from __future__ import annotations

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

import sympy as sp


IntegroDiffKind = Literal[
    "volterra_convert_to_integral_iteration",
    "convolution_laplace_linear",
    "degenerate_kernel_augmented_ode_system",
    "numerical_memory_stepping_template",
    "fredholm_integro_differential_template",
]


def solve_integro_differential_equations(
    *,
    kind: IntegroDiffKind,
    # symbols
    t_symbol: str = "t",
    x_symbol: str = "x",
    s_symbol: str = "s",
    # interval / IC
    t0_: Union[int, float, str, sp.Expr] = 0,
    a_: Union[int, float, str, sp.Expr] = 0,
    b_: Union[int, float, str, sp.Expr] = 1,
    y0_: Union[int, float, str, sp.Expr] = 0,
    # parameters
    a_param: Union[int, float, str, sp.Expr] = 0,   # for y' + a y = ...
    b_param: Union[int, float, str, sp.Expr] = 0,   # optional (not always used)
    lam_: Union[int, float, str, sp.Expr] = 1,      # use numeric by default (avoid "lambda" keyword issues)
    # display strings (kept generic)
    f_str: str = "f(t,y(t))",
    g_str: str = "g(t)",
    kernel_str: str = "K(t,s)",
    conv_k_str: str = "k(t)",
    # iteration depth
    n_iter: int = 3,
    # degenerate kernel representation K(t,s)=sum a_j(t) b_j(s)
    degenerate_a_list: Optional[List[Union[str, sp.Expr]]] = None,
    degenerate_b_list: Optional[List[Union[str, sp.Expr]]] = None,
    # numerical stepping parameters
    dt_symbol: str = r"\Delta t",
    N_steps: int = 5,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    # -----------------------------
    # 1) Symbols + locals
    # -----------------------------
    t = sp.Symbol(t_symbol, real=True)
    x = sp.Symbol(x_symbol, real=True)
    s = sp.Symbol(s_symbol, real=True)

    lam_sym = sp.Symbol("lambda", real=True)

    base_locals: Dict[str, Any] = {
        t_symbol: t,
        x_symbol: x,
        s_symbol: 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,
        # safe lambda mappings:
        "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:
        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)

    def L(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>"
        )

    # parse values
    t0 = P(t0_)
    a = P(a_)
    b = P(b_)
    y0 = P(y0_)
    a_p = P(a_param)
    b_p = P(b_param)
    lam_v = P(lam_) if not isinstance(lam_, (int, float)) else sp.sympify(lam_)

    if simplify:
        t0 = sp.simplify(t0)
        a = sp.simplify(a)
        b = sp.simplify(b)
        y0 = sp.simplify(y0)
        a_p = sp.simplify(a_p)
        b_p = sp.simplify(b_p)
        lam_v = sp.simplify(lam_v)

    lam_disp = L(lam_v)

    # -----------------------------
    # 2) Build templates by kind
    # -----------------------------
    if kind == "volterra_convert_to_integral_iteration":
        # y'(t)=f(t,y(t)) + ∫_{t0}^{t} K(t,s) y(s) ds
        body = (
            "& \\displaystyle \\text{Volterra integro-differential form:} \\\\"
            "& \\displaystyle y'(t)= " + f_str + " + \\int_{t_0}^{t} " + kernel_str + "\\,y(s)\\,ds,\\quad y(t_0)=y_0 \\\\[10pt]"
            "& \\displaystyle \\textbf{Step 1: Integrate from }t_0\\text{ to }t: \\\\"
            "& \\displaystyle y(t)-y_0 = \\int_{t_0}^{t} " + f_str.replace("t", "\\tau") + "\\,d\\tau"
            " + \\int_{t_0}^{t}\\int_{t_0}^{\\tau} " + kernel_str.replace("(t,s)", "(\\tau,s)") + "\\,y(s)\\,ds\\,d\\tau \\\\[10pt]"
            "& \\displaystyle \\Rightarrow y(t)=y_0 + \\int_{t_0}^{t} " + f_str.replace("t", "\\tau") + "\\,d\\tau"
            " + \\int_{t_0}^{t}\\int_{t_0}^{\\tau} " + kernel_str.replace("(t,s)", "(\\tau,s)") + "\\,y(s)\\,ds\\,d\\tau \\\\[10pt]"
            "& \\displaystyle \\textbf{Step 2: Picard-like iteration (successive approximations):} \\\\"
            "& \\displaystyle y_0(t)=y_0 \\;\\text{(or use an ODE-only guess)} \\\\"
        )
        iters = max(1, int(n_iter))
        for k in range(0, iters):
            body += (
                "& \\displaystyle y_{" + str(k + 1) + "}(t)=y_0"
                " + \\int_{t_0}^{t} " + f_str.replace("y(t)", "y_" + str(k) + "(\\tau)").replace("t", "\\tau") + "\\,d\\tau"
                " + \\int_{t_0}^{t}\\int_{t_0}^{\\tau} " + kernel_str.replace("(t,s)", "(\\tau,s)") + "\\,y_" + str(k) + "(s)\\,ds\\,d\\tau \\\\"
            )

        soln_s = card(
            "Integro-Differential Equations — Volterra (Convert to Integral + Iteration)",
            body,
            "This is the standard “method of successive approximations” after converting the integro-differential equation to a Volterra integral equation. Convergence depends on kernel/forcing size.",
        )
        problem_s = "Solve a Volterra integro-differential equation by converting to an integral equation and iterating."
        answer_s = (
            "\\(y(t)=y_0+\\int_{t_0}^{t} f(\\tau,y(\\tau))d\\tau"
            "+\\int_{t_0}^{t}\\int_{t_0}^{\\tau}K(\\tau,s)y(s)ds\\,d\\tau\\)"
            " with Picard iteration to approximate \\(y\\)."
        )
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    if kind == "convolution_laplace_linear":
        # y'(t)+a y(t) = g(t) + ∫_{t0}^{t} k(t-s) y(s) ds  (convolution on [t0,∞))
        # Laplace: (sY - y0) + aY = G + K Y -> Y = (G + y0)/(s+a-K)
        body = (
            "& \\displaystyle \\text{Linear convolution integro-differential:} \\\\"
            "& \\displaystyle y'(t)+" + L(a_p) + "y(t)= " + g_str + " + \\int_{t_0}^{t} "
            + conv_k_str + "(t-\\tau)\\,y(\\tau)\\,d\\tau,\\quad y(t_0)=y_0 \\\\[10pt]"
            "& \\displaystyle \\textbf{Laplace transform (shifted to }t_0=0\\text{ for the template):} \\\\"
            "& \\displaystyle \\mathcal{L}\\{y'\\}=sY(s)-y(0),\\quad \\mathcal{L}\\{(k*y)\\}=K(s)Y(s) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow \\big(sY(s)-y(0)\\big)+" + L(a_p) + "Y(s)=G(s)+K(s)Y(s) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow Y(s)\\,\\big(s+" + L(a_p) + "-K(s)\\big)=G(s)+y(0) \\\\[6pt]"
            "& \\displaystyle \\Rightarrow Y(s)=\\frac{G(s)+y(0)}{s+" + L(a_p) + "-K(s)} \\\\[10pt]"
            "& \\displaystyle \\Rightarrow y(t)=\\mathcal{L}^{-1}\\left\\{\\frac{G(s)+y(0)}{s+" + L(a_p) + "-K(s)}\\right\\}. \\\\"
        )
        soln_s = card(
            "Integro-Differential Equations — Convolution Kernel via Laplace Transform",
            body,
            "Best closed-form path when the kernel is convolution-type and the equation is linear. After obtaining Y(s), invert via partial fractions / known transforms (or leave as inverse Laplace).",
        )
        problem_s = "Solve a linear convolution integro-differential equation using Laplace transforms (template)."
        answer_s = "\\(Y(s)=\\dfrac{G(s)+y(0)}{s+a-K(s)}\\Rightarrow y(t)=\\mathcal{L}^{-1}\\{(G(s)+y(0))/(s+a-K(s))\\}.\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    if kind == "degenerate_kernel_augmented_ode_system":
        # y'(t)+a y(t)=g(t) + ∫_{t0}^{t} Σ a_j(t) b_j(s) y(s) ds
        # Define c_j(t)=∫_{t0}^{t} b_j(s) y(s) ds => c_j'(t)=b_j(t)y(t)
        # Convert to ODE system in (y, c_1..c_N)
        if not degenerate_a_list or not degenerate_b_list or (len(degenerate_a_list) != len(degenerate_b_list)):
            degenerate_a_list = ["a1(t)", "a2(t)"]
            degenerate_b_list = ["b1(s)", "b2(s)"]
        N = len(degenerate_a_list)

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

        body = (
            "& \\displaystyle \\text{Degenerate kernel model: }K(t,s)=\\sum_{j=1}^{N} a_j(t)b_j(s) \\\\"
            "& \\displaystyle K(t,s)= " + kernel_sum + "\\quad (N=" + str(N) + ") \\\\[10pt]"
            "& \\displaystyle \\text{Example structure: } y'(t)+" + L(a_p) + "y(t)= " + g_str + " + \\int_{t_0}^{t}K(t,s)y(s)ds \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\int_{t_0}^{t}K(t,s)y(s)ds = \\sum_{j=1}^{N} a_j(t)\\,c_j(t) \\\\"
            "& \\displaystyle \\text{where } c_j(t)=\\int_{t_0}^{t} b_j(s)\\,y(s)\\,ds \\\\[10pt]"
            "& \\displaystyle \\textbf{Differentiate }c_j(t):\\quad c_j'(t)=b_j(t)\\,y(t),\\qquad c_j(t_0)=0 \\\\[10pt]"
            "& \\displaystyle \\Rightarrow \\textbf{Augmented ODE system in }(y,c_1,\\dots,c_N): \\\\"
            "& \\displaystyle \\begin{cases}"
            "y'(t)= -(" + L(a_p) + ")y(t) + " + g_str + " + \\sum_{j=1}^{N} a_j(t)c_j(t), \\\\"
            "c_1'(t)=b_1(t)y(t),\\;\\dots,\\;c_N'(t)=b_N(t)y(t)"
            "\\end{cases} \\\\[8pt]"
            "& \\displaystyle \\Rightarrow \\text{Now apply standard ODE solvers (RK4/RK45/implicit if stiff).} \\\\"
        )

        soln_s = card(
            "Integro-Differential Equations — Degenerate Kernel → Augmented ODE System",
            body,
            "This is very app-friendly: it turns the integro-differential equation into a finite-dimensional ODE system by introducing memory states c_j(t).",
        )
        problem_s = "Reduce an integro-differential equation with a degenerate kernel to an augmented ODE system."
        answer_s = "\\(c_j(t)=\\int_{t_0}^{t}b_j(s)y(s)ds\\Rightarrow c_j'(t)=b_j(t)y(t)\\), giving an augmented ODE system for \\(y,c_1,\\dots,c_N\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    if kind == "numerical_memory_stepping_template":
        # Generic memory stepping:
        # I_n = ∫_{t0}^{t_n} K(t_n,s) y(s) ds ≈ Σ w_{n,m} K(t_n,t_m) y_m
        # then combine with an ODE step (Euler shown)
        Nn = max(3, int(N_steps))
        body = (
            "& \\displaystyle \\text{Generic Volterra integro-differential (memory):} \\\\"
            "& \\displaystyle y'(t)=f(t,y(t)) + \\int_{t_0}^{t} " + kernel_str + "\\,y(s)\\,ds \\\\[10pt]"
            "& \\displaystyle \\textbf{Time grid: } t_n=t_0+n" + dt_symbol + ",\\quad y_n\\approx y(t_n),\\quad n=0,1,\\dots,N \\\\[6pt]"
            "& \\displaystyle \\textbf{History-sum quadrature: } \\int_{t_0}^{t_n}K(t_n,s)y(s)ds \\approx \\sum_{m=0}^{n} w_{n,m}\\,K(t_n,t_m)\\,y_m \\\\[10pt]"
            "& \\displaystyle \\textbf{Example explicit Euler step:} \\\\"
            "& \\displaystyle \\frac{y_{n+1}-y_n}{" + dt_symbol + "} \\approx f(t_n,y_n)+\\sum_{m=0}^{n} w_{n,m}\\,K(t_n,t_m)\\,y_m \\\\[8pt]"
            "& \\displaystyle \\Rightarrow y_{n+1}=y_n+" + dt_symbol + "\\left[f(t_n,y_n)+\\sum_{m=0}^{n} w_{n,m}\\,K(t_n,t_m)\\,y_m\\right] \\\\[10pt]"
            "& \\displaystyle \\text{Cost note: direct history sum is }O(N^2)\\text{ (each step uses all past values).} \\\\"
        )

        soln_s = card(
            "Integro-Differential Equations — Numerical Memory Stepping (History Sum)",
            body,
            "This is the general-purpose numeric path. For convolution kernels, FFT acceleration can reduce cost (advanced). For stability, implicit time-stepping may be needed.",
        )
        problem_s = "Show a finite-difference / quadrature stepping template for a Volterra integro-differential equation."
        answer_s = "\\(y_{n+1}=y_n+\\Delta t\\,[f(t_n,y_n)+\\sum_{m=0}^{n} w_{n,m}K(t_n,t_m)y_m]\\) (history-sum)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    if kind == "fredholm_integro_differential_template":
        # y'(x)+p(x)y(x)=f(x)+λ∫_a^b K(x,s)y(s)ds
        body = (
            "& \\displaystyle \\text{Fredholm integro-differential (fixed limits):} \\\\"
            "& \\displaystyle y'(x)+p(x)y(x)=f(x)+" + lam_disp + "\\int_{a}^{b} " + kernel_str.replace("(t,s)", "(x,s)") + "\\,y(s)\\,ds \\\\[10pt]"
            "& \\displaystyle \\textbf{Two standard solver routes:} \\\\[2pt]"
            "& \\displaystyle (1)\\ \\text{Quadrature discretization: choose nodes }x_i\\text{ and approximate the integral }\\Rightarrow\\text{ linear system + ODE coupling} \\\\"
            "& \\displaystyle (2)\\ \\text{If kernel is degenerate }K(x,s)=\\sum a_j(x)b_j(s),\\text{ reduce to augmented ODE system (Type: degenerate kernel)} \\\\"
        )
        soln_s = card(
            "Integro-Differential Equations — Fredholm (Template)",
            body,
            "Fredholm fixed-limit coupling makes it non-causal; discretization typically yields a dense system. Degenerate kernels reduce to finite-dimensional coupling.",
        )
        problem_s = "Provide a solution template for a Fredholm integro-differential equation."
        answer_s = "\\(y'(x)+p(x)y(x)=f(x)+\\lambda\\int_a^b K(x,s)y(s)ds\\) (solve via discretization or degenerate-kernel reduction)."
        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_convert_to_integral_iteration', "
        "'convolution_laplace_linear', "
        "'degenerate_kernel_augmented_ode_system', "
        "'numerical_memory_stepping_template', "
        "'fredholm_integro_differential_template'.".format(kind)
    )
