# ============================================================
# File: _solve_nonclassical_differential_equations.py
# Topic: Non-classical Differential Equations
#
# Covers (concept + core algorithm template):
#   1) Delay Differential Equations (DDE): method of steps + history function
#   2) Stochastic Differential Equations (SDE): Euler–Maruyama / Milstein
#   3) Fractional Differential Equations (FDE): Caputo derivative + history weights
#
# This solver is "app-friendly": it produces a structured explanation and
# canonical update formulas, not a full simulation engine.
#
# Output:
#   { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
#
# Notes:
#   - Avoid f-strings to prevent LaTeX brace errors in HTML strings.
#   - You can later plug these templates into actual numerical simulators.
# ============================================================

from __future__ import annotations

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

import sympy as sp


NonclassicalKind = Literal[
    "dde_method_of_steps",
    "sde_euler_maruyama",
    "sde_milstein",
    "fractional_caputo_l1_scheme",
]


def solve_nonclassical_differential_equations(
    *,
    kind: NonclassicalKind,
    # optional display parameters (kept symbolic)
    tau_: Union[int, float, str, sp.Expr] = "1",
    alpha_: Union[int, float, str, sp.Expr] = "0.5",
    dt_: Union[int, float, str, sp.Expr] = "0.01",
    # symbols
    t_symbol: str = "t",
    y_symbol: str = "y",
    x_symbol: str = "X",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    t = sp.Symbol(t_symbol, real=True)
    tau = sp.Symbol("tau", positive=True, real=True)
    alpha = sp.Symbol("alpha", positive=True, real=True)
    dt = sp.Symbol("Delta t", positive=True, real=True)

    base_locals = {
        t_symbol: t,
        "tau": tau,
        "alpha": alpha,
        "dt": dt,
        "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,
    }
    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>"
        )

    tau_v = P(tau_)
    alpha_v = P(alpha_)
    dt_v = P(dt_)
    if simplify:
        tau_v = sp.simplify(tau_v)
        alpha_v = sp.simplify(alpha_v)
        dt_v = sp.simplify(dt_v)

    # ------------------------------------------------------------
    # DDE: method of steps
    # ------------------------------------------------------------
    if kind == "dde_method_of_steps":
        body = (
            "& \\displaystyle \\text{Delay DE (DDE): } y'(t)=f\\big(t,\\,y(t),\\,y(t-\\tau)\\big) \\\\[6pt]"
            "& \\displaystyle \\text{Needs history: } y(t)=\\phi(t)\\quad \\text{for } t\\in[t_0-\\tau,\\,t_0] \\\\[10pt]"
            "& \\displaystyle \\textbf{Method of steps:} \\\\[2pt]"
            "& \\displaystyle \\text{On }[t_0,\\,t_0+\\tau]:\\; y(t-\\tau)=\\phi(t-\\tau)\\text{ is known }\\Rightarrow\\text{ solve an ODE} \\\\[6pt]"
            "& \\displaystyle y'(t)=f\\big(t,\\,y(t),\\,\\phi(t-\\tau)\\big),\\quad t\\in[t_0,\\,t_0+\\tau] \\\\[10pt]"
            "& \\displaystyle \\text{Then reuse computed }y\\text{ as history on }[t_0+\\tau,\\,t_0+2\\tau],\\;\\text{etc.} \\\\"
        )
        soln_s = card(
            "Non-classical DE — Delay Differential Equation (Method of Steps)",
            body,
            "DDE IVPs require a whole function history on an interval, not just y(t0). Each step interval turns the DDE into an ODE because the delayed term is already known there.",
        )
        problem_s = "Explain how to solve a DDE using the method of steps."
        answer_s = "\\(\\text{Provide history }\\phi\\text{ on }[t_0-\\tau,t_0]\\text{ and solve interval-by-interval as ODEs.}\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # SDE: Euler–Maruyama
    # ------------------------------------------------------------
    if kind == "sde_euler_maruyama":
        body = (
            "& \\displaystyle \\text{Ito SDE: } dX_t=a(X_t,t)\\,dt + b(X_t,t)\\,dW_t \\\\[8pt]"
            "& \\displaystyle \\text{Discretize } t_n=t_0+n\\Delta t,\\; \\Delta W_n \\sim \\mathcal{N}(0,\\Delta t) \\\\[8pt]"
            "& \\displaystyle \\textbf{Euler\\text{–}Maruyama:} \\\\[2pt]"
            "& \\displaystyle X_{n+1}=X_n + a(X_n,t_n)\\,\\Delta t + b(X_n,t_n)\\,\\Delta W_n \\\\[10pt]"
            "& \\displaystyle \\text{Key: solution is a random process; one computes sample paths and statistics.} \\\\"
        )
        soln_s = card(
            "Non-classical DE — SDE (Euler–Maruyama)",
            body,
            "For SDEs, outputs are random sample paths. A single run gives one path; repeated runs estimate mean/variance/distributions.",
        )
        problem_s = "State the Euler–Maruyama scheme for an Ito SDE."
        answer_s = "\\(X_{n+1}=X_n+a(X_n,t_n)\\Delta t+b(X_n,t_n)\\Delta W_n,\\;\\Delta W_n\\sim\\mathcal{N}(0,\\Delta t).\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # SDE: Milstein
    # ------------------------------------------------------------
    if kind == "sde_milstein":
        body = (
            "& \\displaystyle \\text{Ito SDE: } dX_t=a(X_t,t)\\,dt + b(X_t,t)\\,dW_t \\\\[8pt]"
            "& \\displaystyle \\textbf{Milstein (1D):} \\\\[2pt]"
            "& \\displaystyle X_{n+1}=X_n + a_n\\,\\Delta t + b_n\\,\\Delta W_n + \\frac12 b_n\\,b_x(X_n,t_n)\\big((\\Delta W_n)^2-\\Delta t\\big) \\\\[6pt]"
            "& \\displaystyle \\text{where } a_n=a(X_n,t_n),\\; b_n=b(X_n,t_n),\\; b_x=\\frac{\\partial b}{\\partial x} \\\\[10pt]"
            "& \\displaystyle \\text{Improves strong accuracy when }b\\text{ depends on }X. \\\\"
        )
        soln_s = card(
            "Non-classical DE — SDE (Milstein Correction)",
            body,
            "Milstein adds a correction term involving b_x and (ΔW^2−Δt), improving pathwise (strong) accuracy compared to Euler–Maruyama.",
        )
        problem_s = "Write the Milstein scheme for a 1D Ito SDE."
        answer_s = "\\(X_{n+1}=X_n+a_n\\Delta t+b_n\\Delta W_n+\\tfrac12 b_n b_x\\big((\\Delta W_n)^2-\\Delta t\\big).\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # ------------------------------------------------------------
    # Fractional DE: Caputo derivative, L1 scheme (history weights)
    # ------------------------------------------------------------
    if kind == "fractional_caputo_l1_scheme":
        # L1 scheme sketch: D_c^α y(t_{n}) ≈ (1/Γ(2-α)) * Σ_{k=0}^{n-1} (y_{k+1}-y_k)/Δt * [ (n-k)^{1-α} - (n-k-1)^{1-α} ]
        # We'll present it as a weighted history convolution.
        Gamma = sp.gamma
        body = (
            "& \\displaystyle \\text{Fractional DE (Caputo, }0<\\alpha<1\\text{): } {}^{C}D_t^{\\alpha} y(t)=f(t,y) \\\\[8pt]"
            "& \\displaystyle \\text{Key feature: memory (depends on whole past).} \\\\[10pt]"
            "& \\displaystyle \\textbf{L1 discretization (uniform }\\Delta t\\text{):} \\\\[2pt]"
            "& \\displaystyle {}^{C}D_t^{\\alpha} y(t_n) \\approx \\frac{1}{\\Gamma(2-\\alpha)}\\sum_{k=0}^{n-1}\\frac{y_{k+1}-y_k}{\\Delta t}\\Big((n-k)^{1-\\alpha}-(n-k-1)^{1-\\alpha}\\Big) \\\\[10pt]"
            "& \\displaystyle \\text{This is a history-weighted sum (convolution-like). Cost can be }O(N^2)\\text{ unless accelerated.} \\\\"
        )
        soln_s = card(
            "Non-classical DE — Fractional (Caputo) via History Weights (L1 Scheme)",
            body,
            "Fractional derivatives are nonlocal. Numerically, the derivative becomes a weighted sum over past increments, which is why standard ODE solvers don’t directly apply.",
        )
        problem_s = "Show a standard discretization idea for a Caputo fractional derivative (0<α<1)."
        answer_s = "Use a history-weighted sum, e.g. the L1 scheme: \\({}^{C}D_t^{\\alpha}y(t_n)\\approx \\frac{1}{\\Gamma(2-\\alpha)}\\sum_{k=0}^{n-1}\\frac{y_{k+1}-y_k}{\\Delta t}((n-k)^{1-\\alpha}-(n-k-1)^{1-\\alpha})\\)."
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError(
        "Unknown kind={!r}. Use one of: "
        "'dde_method_of_steps','sde_euler_maruyama','sde_milstein','fractional_caputo_l1_scheme'.".format(kind)
    )
