# ============================================================
# File: _solve_ode_special_function_families.py
# Topic: Special-Function ODEs (Bessel, Legendre, Hermite, Laguerre, Airy)
#
# We implement an app-friendly "recognize + write canonical solution" solver
# for standard 2nd-order linear ODE families:
#
# 1) Bessel:
#      x^2 y'' + x y' + ( (k^2)x^2 - nu^2 ) y = 0
#    => y = C1*J_nu(kx) + C2*Y_nu(kx)
#
# 2) Modified Bessel:
#      x^2 y'' + x y' - ( (k^2)x^2 + nu^2 ) y = 0
#    => y = C1*I_nu(kx) + C2*K_nu(kx)
#
# 3) Legendre:
#      (1-x^2) y'' - 2x y' + l(l+1) y = 0
#    => y = C1*P_l(x) + C2*Q_l(x)
#
# 4) Hermite (physicists'):
#      y'' - 2x y' + 2n y = 0
#    => y = C1*H_n(x) + C2*(2nd independent solution)
#    (We keep the 2nd solution symbolic for maximum compatibility.)
#
# 5) Laguerre:
#      x y'' + (1-x) y' + n y = 0
#    => y = C1*L_n(x) + C2*U(-n,1,x)
#
# 6) Associated Laguerre:
#      x y'' + (alpha+1-x) y' + n y = 0
#    => y = C1*L_n^alpha(x) + C2*U(-n,alpha+1,x)
#
# 7) Airy:
#      y'' - x y = 0
#    => y = C1*Ai(x) + C2*Bi(x)
#
# Notes:
#   - We template recognized solutions using SymPy special funcs where stable.
#   - Some SymPy builds do NOT expose hyperu; so we keep Tricomi U(a,b,x)
#     as a symbolic function U(a,b,x) to avoid runtime errors.
#   - Optional regularity at x=0 for Bessel families:
#       "regular_at_0" => sets C2=0 (Y_nu or K_nu typically singular at 0)
#   - SymPy does not consistently expose Legendre-Q across versions, so we:
#       * build Q_l(x) by recurrence when l is a nonnegative integer
#       * otherwise keep Q(l,x) symbolic (safe)
# ============================================================

from __future__ import annotations

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

import sympy as sp


def _legendre_Q_integer_degree(ell_int: int, x: sp.Symbol, *, simplify: bool = True) -> sp.Expr:
    """Version-safe Legendre function of the 2nd kind Q_ell(x) for integer ell >= 0.

    SymPy does not consistently expose a built-in `legendre_Q` across versions.
    For the classic Legendre ODE in physics/engineering, ell is typically a
    nonnegative integer (degree). In that case, Q_ell(x) can be generated by
    the standard 3-term recurrence with base cases:

        Q_0(x) = (1/2) * log((1+x)/(1-x))
        Q_1(x) = x*Q_0(x) - 1

    and recurrence:

        (n+1) Q_{n+1}(x) = (2n+1) x Q_n(x) - n Q_{n-1}(x).
    """
    n = int(ell_int)
    if n < 0:
        raise ValueError("ell_int must be >= 0")

    Q0 = sp.Rational(1, 2) * sp.log((1 + x) / (1 - x))
    if n == 0:
        return sp.simplify(Q0) if simplify else Q0

    Q1 = x * Q0 - 1
    if n == 1:
        return sp.simplify(Q1) if simplify else Q1

    Qnm1, Qn = Q0, Q1
    for k in range(1, n):
        Qnp1 = ((2 * k + 1) * x * Qn - k * Qnm1) / (k + 1)
        if simplify:
            Qnp1 = sp.simplify(Qnp1)
        Qnm1, Qn = Qn, Qnp1

    return Qn


SpecialOdeKind = Literal[
    "bessel",
    "modified_bessel",
    "legendre",
    "hermite",
    "laguerre",
    "assoc_laguerre",
    "airy",
]


def solve_ode_special_function_families(
    *,
    kind: SpecialOdeKind,
    x_symbol: str = "x",
    y_symbol: str = "y",
    C1_symbol: str = "C_1",
    C2_symbol: str = "C_2",
    # Parameters (used depending on kind)
    nu: Union[int, float, str, sp.Expr] = 0,        # Bessel order
    k: Union[int, float, str, sp.Expr] = 1,         # scaling in argument kx
    ell: Union[int, float, str, sp.Expr] = 0,       # Legendre degree
    n: Union[int, float, str, sp.Expr] = 0,         # Hermite/Laguerre integer
    alpha: Union[int, float, str, sp.Expr] = 0,     # assoc Laguerre
    # Optional simple condition
    regular_at_0: bool = False,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Returns:
        { "problem_s": "...", "soln_s": "...", "answer_s": "..." }

    Output style:
      - HTML with \\(\\displaystyle ...\\) and aligned steps
      - answer_s is the final solution statement
    """
    x = sp.Symbol(x_symbol, real=True)
    y = sp.Function(y_symbol)
    yx = y(x)

    C1 = sp.Symbol(C1_symbol)
    C2 = sp.Symbol(C2_symbol)

    # Symbolic Tricomi U(a,b,x) (safe across SymPy versions)
    Ufun = sp.Function("U")

    base_locals = {
        x_symbol: x,
        y_symbol: y,
        "pi": sp.pi,
        "E": sp.E,
        "exp": sp.exp,
        "log": sp.log,
        "ln": sp.log,
        "sqrt": sp.sqrt,
        "sin": sp.sin,
        "cos": sp.cos,
        "tan": sp.tan,
        "sinh": sp.sinh,
        "cosh": sp.cosh,
        "tanh": sp.tanh,
        "Abs": sp.Abs,
        "besselj": sp.besselj,
        "bessely": sp.bessely,
        "besseli": sp.besseli,
        "besselk": sp.besselk,
        "legendre": sp.legendre,
        "assoc_laguerre": sp.assoc_laguerre,
        "laguerre": sp.laguerre,
        "hermite": sp.hermite,
        "airyai": sp.airyai,
        "airybi": sp.airybi,
        "hyper": sp.hyper,
        "U": Ufun,
    }
    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 L(expr: sp.Expr) -> str:
        return sp.latex(expr)

    Nu = P(nu)
    K = P(k)
    Ell = P(ell)
    N = P(n)
    A = P(alpha)

    if simplify:
        Nu = sp.simplify(Nu)
        K = sp.simplify(K)
        Ell = sp.simplify(Ell)
        N = sp.simplify(N)
        A = sp.simplify(A)

    # Build canonical ODE + solution template
    ode_ltx = ""
    sol_expr = None
    family_name = ""
    steps = []

    if kind == "bessel":
        family_name = "Bessel"
        ode_ltx = (
            f"{L(x)}^2 y'' + {L(x)} y' + \\left(({L(K)}^2 {L(x)}^2 - {L(Nu)}^2)\\right) y = 0"
        )
        arg = sp.simplify(K * x) if simplify else (K * x)
        sol_expr = C1 * sp.besselj(Nu, arg) + C2 * sp.bessely(Nu, arg)

        steps.append("Recognize the Bessel form.")
        steps.append("Use the canonical solution in terms of \\(J_\\nu\\) and \\(Y_\\nu\\).")

        if regular_at_0:
            sol_expr = sp.simplify(C1 * sp.besselj(Nu, arg)) if simplify else (C1 * sp.besselj(Nu, arg))
            steps.append("Apply regularity at \\(x=0\\): discard \\(Y_\\nu\\) (singular), so \\(C_2=0\\).")

    elif kind == "modified_bessel":
        family_name = "Modified Bessel"
        ode_ltx = (
            f"{L(x)}^2 y'' + {L(x)} y' - \\left(({L(K)}^2 {L(x)}^2 + {L(Nu)}^2)\\right) y = 0"
        )
        arg = sp.simplify(K * x) if simplify else (K * x)
        sol_expr = C1 * sp.besseli(Nu, arg) + C2 * sp.besselk(Nu, arg)

        steps.append("Recognize the modified Bessel form.")
        steps.append("Use the canonical solution in terms of \\(I_\\nu\\) and \\(K_\\nu\\).")

        if regular_at_0:
            sol_expr = sp.simplify(C1 * sp.besseli(Nu, arg)) if simplify else (C1 * sp.besseli(Nu, arg))
            steps.append("Apply regularity at \\(x=0\\): discard \\(K_\\nu\\) (singular), so \\(C_2=0\\).")

    elif kind == "legendre":
        family_name = "Legendre"
        ode_ltx = (
            f"(1-{L(x)}^2) y'' - 2{L(x)} y' + {L(Ell)}({L(Ell)}+1) y = 0"
        )

        # SymPy does not reliably expose a built-in Legendre-Q across versions.
        # Best, version-safe behavior for an educational app:
        #   - If ell is a nonnegative integer, build Q_ell(x) explicitly via recurrence
        #   - Otherwise keep Q_ell(x) symbolic as Q(ell, x)
        P_l = sp.legendre(Ell, x)

        Q_l: sp.Expr
        if Ell.is_number and bool(Ell.is_integer):
            try:
                ell_int = int(Ell)
                if ell_int >= 0:
                    Q_l = _legendre_Q_integer_degree(ell_int, x, simplify=simplify)
                    steps.append("Since \\(\\ell\\) is an integer, generate \\(Q_\\ell(x)\\) via the 3-term recurrence (version-safe).")
                else:
                    Q_l = sp.Function("Q")(Ell, x)
                    steps.append("For negative/nonstandard degree, keep \\(Q_\\ell\\) symbolic as \\(Q(\\ell,x)\\).")
            except Exception:
                Q_l = sp.Function("Q")(Ell, x)
                steps.append("Keep \\(Q_\\ell\\) symbolic as \\(Q(\\ell,x)\\) (could not expand).")
        else:
            Q_l = sp.Function("Q")(Ell, x)
            steps.append("Keep \\(Q_\\ell\\) symbolic as \\(Q(\\ell,x)\\) (non-integer degree).")

        sol_expr = C1 * P_l + C2 * Q_l

        steps.append("Recognize the Legendre form.")
        steps.append("Write the general solution using \\(P_\\ell(x)\\) and \\(Q_\\ell(x)\\).")

    elif kind == "hermite":
        family_name = "Hermite"
        ode_ltx = f"y'' - 2{L(x)} y' + 2({L(N)}) y = 0"

        sol1 = sp.hermite(N, x)
        sol2 = sp.Function("H2")(x)  # keep symbolic (version-safe)

        sol_expr = C1 * sol1 + C2 * sol2
        steps.append("Recognize the Hermite form.")
        steps.append("For integer \\(n\\ge 0\\), one solution is the Hermite polynomial \\(H_n(x)\\).")
        steps.append("A second independent solution exists (often hypergeometric); kept symbolic as \\(H_2(x)\\).")

    elif kind == "laguerre":
        family_name = "Laguerre"
        ode_ltx = f"{L(x)} y'' + (1-{L(x)}) y' + {L(N)} y = 0"

        sol1 = sp.laguerre(N, x)
        sol2 = Ufun(-N, 1, x)  # Tricomi U(a,b,x) kept symbolic

        sol_expr = C1 * sol1 + C2 * sol2
        steps.append("Recognize the Laguerre form.")
        steps.append("Write one solution as \\(L_n(x)\\).")
        steps.append("Second independent solution is commonly written using Tricomi's \\(U(a,b,x)\\); kept symbolic as \\(U(-n,1,x)\\).")

    elif kind == "assoc_laguerre":
        family_name = "Associated Laguerre"
        ode_ltx = f"{L(x)} y'' + ({L(A)}+1-{L(x)}) y' + {L(N)} y = 0"

        sol1 = sp.assoc_laguerre(N, A, x)
        sol2 = Ufun(-N, A + 1, x)  # Tricomi U(a,b,x) kept symbolic

        sol_expr = C1 * sol1 + C2 * sol2
        steps.append("Recognize the associated Laguerre form.")
        steps.append("Write one solution as \\(L_n^{\\alpha}(x)\\).")
        steps.append("Second independent solution is commonly written using Tricomi's \\(U(a,b,x)\\); kept symbolic as \\(U(-n,\\alpha+1,x)\\).")

    elif kind == "airy":
        family_name = "Airy"
        ode_ltx = f"y'' - {L(x)} y = 0"

        sol_expr = C1 * sp.airyai(x) + C2 * sp.airybi(x)
        steps.append("Recognize the Airy form.")
        steps.append("Write the solution using \\(\\operatorname{Ai}(x)\\) and \\(\\operatorname{Bi}(x)\\).")

    else:
        raise ValueError(f"Unknown kind={kind!r}.")

    if simplify and sol_expr is not None:
        sol_expr = sp.simplify(sol_expr)

    # Build HTML safely (avoid f-string braces in LaTeX by using concatenation)
    step_lines = ""
    for s in steps:
        step_lines += "& \\displaystyle \\Rightarrow \\text{" + s + "} \\\\"

    problem_s = "Solve the special-function ODE: \\( " + ode_ltx + " \\)"

    soln_s = (
        "<div class='wrap'><div class='card'>"
        "<div class='badge'>Special-Function ODE — " + family_name + "<br/></div>"
        "\\(\\displaystyle \\begin{aligned}"
        "& \\displaystyle " + ode_ltx + " \\\\[8pt]"
        + step_lines +
        "& \\displaystyle \\Rightarrow \\text{General solution: } " + L(yx) + " = " + L(sol_expr) + " \\\\"
        "\\end{aligned}\\)"
        "</div></div>"
    )

    answer_s = "\\( " + L(yx) + " = " + L(sol_expr) + " \\)"

    return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}
