# ============================================================
# File: _solve_ode_series_methods_power_frobenius.py
# Topic: Series Methods for ODEs (Power Series + Frobenius)
#
# We support two classic series methods:
#
# (A) Power Series about an ordinary point x0:
#       y'' + P(x) y' + Q(x) y = R(x)
#     where P,Q,R are analytic at x0.
#     We build a truncated series:
#       y ≈ Σ_{n=0..N} a_n (x-x0)^n
#     and solve for a2..aN in terms of a0,a1 (or in terms of ICs).
#
# (B) Frobenius about a regular singular point x0:
#       (x-x0)^2 y'' + (x-x0) P(x) y' + Q(x) y = 0
#     where P,Q analytic at x0.
#     Assume:
#       y = Σ_{n=0..N} a_n (x-x0)^{n+r}, a0≠0
#     Indicial equation:
#       r(r-1) + p0 r + q0 = 0,  where p0=P(x0), q0=Q(x0)
#     Recurrence:
#       a_n = - [ Σ_{k=1..n} ( p_k (n-k+r) + q_k ) a_{n-k} ] / D_n
#       D_n = (n+r)(n+r-1) + p0(n+r) + q0
#
# Notes:
#   - Output is truncated (order N) to be app-friendly.
#   - If SymPy can’t series-expand, it keeps series/integrals symbolic.
# ============================================================

from __future__ import annotations

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

import sympy as sp


SeriesMethodKind = Literal["power_series", "frobenius"]


def solve_ode_series_methods_power_frobenius(
    *,
    method: SeriesMethodKind,
    # --- Ordinary-point form ---
    P_of_x: Union[str, sp.Expr, None] = None,
    Q_of_x: Union[str, sp.Expr, None] = None,
    R_of_x: Union[str, sp.Expr, None] = None,
    # --- Frobenius regular-singular form ---
    P_reg_of_x: Union[str, sp.Expr, None] = None,  # in (x-x0)^2 y'' + (x-x0) P(x) y' + Q(x) y = 0
    Q_reg_of_x: Union[str, sp.Expr, None] = None,
    # --- common controls ---
    x_symbol: str = "x",
    y_symbol: str = "y",
    x0: Union[int, float, str, sp.Expr] = 0,
    N: int = 6,  # truncation order
    # ICs for power series:
    y0: Optional[Union[int, float, str, sp.Expr]] = None,    # y(x0)
    yp0: Optional[Union[int, float, str, sp.Expr]] = None,   # y'(x0)
    # Frobenius choices:
    pick_root_index: int = 0,  # 0 => r1, 1 => r2 (if distinct)
    a0_value: Union[int, float, str, sp.Expr] = 1,  # choose a0 for Frobenius series
    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 truncated series form.
    """
    if N < 2:
        raise ValueError("N should be >= 2 for meaningful series output.")

    x = sp.Symbol(x_symbol, real=True)
    y = sp.Function(y_symbol)
    yx = y(x)

    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,
        "asin": sp.asin,
        "acos": sp.asin,
        "atan": sp.atan,
        "sinh": sp.sinh,
        "cosh": sp.cosh,
        "tanh": sp.tanh,
        "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, None]) -> Optional[sp.Expr]:
        if v is None:
            return None
        return sp.sympify(v, locals=base_locals)

    def L(expr: sp.Expr) -> str:
        return sp.latex(expr)

    x0e = sp.sympify(x0, locals=base_locals)
    if simplify:
        x0e = sp.simplify(x0e)

    h = sp.Symbol("h", real=True)

    # -----------------------------
    # Helper: make truncated series of an expression about x0
    # -----------------------------
    def series_coeffs_about_x0(expr_x: sp.Expr, order: int) -> List[sp.Expr]:
        """
        Return [c0,c1,...,c_order] where:
          expr(x0+h) = Σ_{k=0..order} c_k h^k + O(h^{order+1})
        """
        expr_h = expr_x.subs(x, x0e + h)
        ser = sp.series(expr_h, h, 0, order + 1).removeO()
        if simplify:
            ser = sp.simplify(ser)
        poly = sp.Poly(ser, h)
        out = []
        for k in range(order + 1):
            out.append(sp.simplify(poly.coeffs()[poly.monoms().index((k,))]) if (k,) in poly.monoms() else sp.Integer(0))
        return out

    # robust coefficients (avoid index lookup cost/errors)
    def series_coeffs_about_x0_safe(expr_x: sp.Expr, order: int) -> List[sp.Expr]:
        expr_h = expr_x.subs(x, x0e + h)
        ser = sp.series(expr_h, h, 0, order + 1).removeO()
        if simplify:
            ser = sp.simplify(ser)
        poly = sp.Poly(ser, h, domain="EX")
        out = [sp.Integer(0)] * (order + 1)
        for (k,), c in poly.terms():
            if 0 <= k <= order:
                out[k] = sp.simplify(c) if simplify else c
        return out

    # -----------------------------
    # (A) Power Series (ordinary point)
    # -----------------------------
    if method == "power_series":
        if P_of_x is None or Q_of_x is None or R_of_x is None:
            raise ValueError("For method='power_series', provide P_of_x, Q_of_x, R_of_x for y'' + P y' + Q y = R.")

        Pexpr = P(P_of_x)
        Qexpr = P(Q_of_x)
        Rexpr = P(R_of_x)

        if simplify:
            Pexpr = sp.simplify(Pexpr)
            Qexpr = sp.simplify(Qexpr)
            Rexpr = sp.simplify(Rexpr)

        # Build truncated y series: y = Σ a_n h^n
        a = [sp.Symbol(f"a_{n}") for n in range(N + 1)]
        y_ser = sp.Add(*[a[n] * h**n for n in range(N + 1)])
        y1_ser = sp.diff(y_ser, h)
        y2_ser = sp.diff(y1_ser, h)

        # Expand P,Q,R about x0: in h
        Pser = sp.series(Pexpr.subs(x, x0e + h), h, 0, N + 1).removeO()
        Qser = sp.series(Qexpr.subs(x, x0e + h), h, 0, N + 1).removeO()
        Rser = sp.series(Rexpr.subs(x, x0e + h), h, 0, N + 1).removeO()
        if simplify:
            Pser = sp.simplify(Pser)
            Qser = sp.simplify(Qser)
            Rser = sp.simplify(Rser)

        # Substitute into: y'' + P y' + Q y - R = 0, collect up to h^N
        expr = y2_ser + Pser * y1_ser + Qser * y_ser - Rser
        expr_trunc = sp.series(expr, h, 0, N + 1).removeO()
        if simplify:
            expr_trunc = sp.simplify(expr_trunc)

        poly = sp.Poly(expr_trunc, h, domain="EX")
        eqs = []
        for n in range(N + 1):
            cn = poly.coeff_monomial(h**n)
            if simplify:
                cn = sp.simplify(cn)
            eqs.append(sp.Eq(cn, 0))

        # Unknowns: a2..aN (a0,a1 are free unless IC given)
        unknowns = a[2:]
        sol_map = {}
        if unknowns:
            try:
                sol_list = sp.solve([e.lhs for e in eqs], unknowns, dict=True)
                if sol_list:
                    sol_map = sol_list[0]
            except Exception:
                sol_map = {}

        # Apply IC if present: y(x0)=y0 => a0=y0 ; y'(x0)=yp0 => a1=yp0
        use_ic = (y0 is not None) and (yp0 is not None)
        if use_ic:
            y0e = sp.sympify(y0, locals=base_locals)
            yp0e = sp.sympify(yp0, locals=base_locals)
            if simplify:
                y0e = sp.simplify(y0e)
                yp0e = sp.simplify(yp0e)
            sol_map = dict(sol_map)
            sol_map[a[0]] = y0e
            sol_map[a[1]] = yp0e

        # Build final truncated series in x (replace h with x-x0)
        y_final_h = y_ser
        if sol_map:
            y_final_h = y_final_h.subs(sol_map)
        if simplify:
            y_final_h = sp.simplify(y_final_h)

        y_final_x = y_final_h.subs(h, x - x0e)
        if simplify:
            y_final_x = sp.simplify(y_final_x)

        # Build display strings safely (avoid f-strings with LaTeX braces)
        ode_ltx = "y'' + " + L(Pexpr) + "y' + " + L(Qexpr) + "y = " + L(Rexpr)
        center_ltx = L(x0e)

        problem_s = "Solve near \\(x_0=" + center_ltx + "\\): \\( " + ode_ltx + " \\)"

        notes = []
        if not sol_map:
            notes.append("Could not fully solve for coefficients; returning a formal truncated series with free parameters.")
        note_s = " ".join(notes)

        # Show recurrence info by listing first few solved coefficients if available
        shown = []
        for idx in range(2, min(N + 1, 6)):
            if a[idx] in sol_map:
                shown.append("a_" + str(idx) + "=" + L(sol_map[a[idx]]))
        coeff_line = ""
        if shown:
            coeff_line = " & \\displaystyle \\Rightarrow \\text{Coefficients: } " + ",\\; ".join(shown) + " \\\\"

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Series Method (Ordinary Point) — Power Series<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            "& \\displaystyle " + ode_ltx + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow y=\\sum_{n=0}^{\\infty} a_n (x-x_0)^n,\\; x_0=" + center_ltx + " \\\\"
            "& \\displaystyle \\Rightarrow \\text{Truncate to order } " + str(N) + " \\text{ and substitute into the ODE} \\\\"
            "& \\displaystyle \\Rightarrow \\text{Collect powers of } (x-x_0)^n \\text{ and set each coefficient to } 0 \\\\"
            + coeff_line +
            "& \\displaystyle \\Rightarrow \\text{Truncated series solution: } y(x)\\approx " + L(y_final_x) + " \\\\"
            "\\end{aligned}\\)"
            "<div class='note'>" + note_s + "</div>"
            "</div></div>"
        )

        answer_s = "\\( y(x)\\approx " + L(y_final_x) + " \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # -----------------------------
    # (B) Frobenius (regular singular point)
    # -----------------------------
    if method == "frobenius":
        if P_reg_of_x is None or Q_reg_of_x is None:
            raise ValueError("For method='frobenius', provide P_reg_of_x and Q_reg_of_x for (x-x0)^2 y'' + (x-x0)P y' + Q y = 0.")

        Pexpr = P(P_reg_of_x)
        Qexpr = P(Q_reg_of_x)
        if simplify:
            Pexpr = sp.simplify(Pexpr)
            Qexpr = sp.simplify(Qexpr)

        # Get p_k, q_k coefficients of P(x0+h), Q(x0+h) up to N
        p = series_coeffs_about_x0_safe(Pexpr, N)
        q = series_coeffs_about_x0_safe(Qexpr, N)

        p0 = p[0]
        q0 = q[0]

        r = sp.Symbol("r")
        indicial = sp.expand(r * (r - 1) + p0 * r + q0)
        r_roots = sp.solve(sp.Eq(indicial, 0), r)

        # choose r value
        r_val = None
        if isinstance(r_roots, list) and len(r_roots) >= 1:
            idx = 0 if pick_root_index not in (0, 1) else pick_root_index
            idx = min(idx, len(r_roots) - 1)
            r_val = sp.simplify(r_roots[idx]) if simplify else r_roots[idx]
        else:
            # fallback symbolic root
            r_val = sp.Symbol("r_0")

        a0e = sp.sympify(a0_value, locals=base_locals)
        if simplify:
            a0e = sp.simplify(a0e)

        # Build coefficients a_n by recurrence
        a = [sp.Symbol(f"a_{n}") for n in range(N + 1)]
        a_vals = {a[0]: a0e}

        def D(n: int) -> sp.Expr:
            return (n + r_val) * (n + r_val - 1) + p0 * (n + r_val) + q0

        for n in range(1, N + 1):
            # numerator: Σ_{k=1..n} ( p_k (n-k+r) + q_k ) a_{n-k}
            num = sp.Integer(0)
            for k in range(1, n + 1):
                pk = p[k] if k <= N else sp.Integer(0)
                qk = q[k] if k <= N else sp.Integer(0)
                num += (pk * (n - k + r_val) + qk) * a_vals[a[n - k]]
            an = -num / D(n)
            if simplify:
                an = sp.simplify(an)
            a_vals[a[n]] = an

        # Build y series: Σ a_n h^{n+r}
        y_series_h = sp.Add(*[a_vals[a[n]] * h ** (n + r_val) for n in range(N + 1)])
        if simplify:
            y_series_h = sp.simplify(y_series_h)

        y_series_x = y_series_h.subs(h, x - x0e)
        if simplify:
            y_series_x = sp.simplify(y_series_x)

        # Display strings
        ode_ltx = "(" + L(x - x0e) + ")^2 y'' + (" + L(x - x0e) + ")(" + L(Pexpr) + ") y' + (" + L(Qexpr) + ") y = 0"
        center_ltx = L(x0e)

        r_roots_ltx = ""
        if isinstance(r_roots, list) and len(r_roots) >= 1:
            r_roots_ltx = ",\\; ".join([L(rr) for rr in r_roots])

        notes = []
        if len(r_roots) == 1:
            notes.append("Repeated root; a second independent solution may involve \\(\\ln(x-x_0)\\).")
        note_s = " ".join(notes)

        # show first few coefficients
        shown = []
        for idx in range(1, min(N + 1, 6)):
            shown.append("a_" + str(idx) + "=" + L(a_vals[a[idx]]))
        coeff_line = " & \\displaystyle \\Rightarrow \\text{Coefficients: } " + ",\\; ".join(shown) + " \\\\"

        problem_s = "Solve near \\(x_0=" + center_ltx + "\\) using Frobenius: \\( " + ode_ltx + " \\)"

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Series Method (Regular Singular Point) — Frobenius<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            "& \\displaystyle " + ode_ltx + " \\\\[8pt]"
            "& \\displaystyle \\Rightarrow y=\\sum_{n=0}^{\\infty} a_n (x-x_0)^{n+r},\\; a_0\\neq 0,\\; x_0=" + center_ltx + " \\\\"
            "& \\displaystyle \\Rightarrow \\text{Indicial equation: } r(r-1)+p_0 r+q_0=0 \\\\"
            "& \\displaystyle \\Rightarrow p_0=P(x_0)=" + L(p0) + ",\\; q_0=Q(x_0)=" + L(q0) + " \\\\"
            "& \\displaystyle \\Rightarrow r(r-1)+" + L(p0) + "r+" + L(q0) + "=0 \\\\"
            + (
                "& \\displaystyle \\Rightarrow r = " + r_roots_ltx + " \\\\"
                if r_roots_ltx else
                "& \\displaystyle \\Rightarrow r = " + L(r_val) + " \\\\"
            ) +
            "& \\displaystyle \\Rightarrow \\text{Choose } r=" + L(r_val) + " \\text{ and compute recurrence} \\\\"
            + coeff_line +
            "& \\displaystyle \\Rightarrow \\text{Truncated Frobenius series: } y(x)\\approx " + L(y_series_x) + " \\\\"
            "\\end{aligned}\\)"
            "<div class='note'>" + note_s + "</div>"
            "</div></div>"
        )

        answer_s = "\\( y(x)\\approx " + L(y_series_x) + " \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    raise ValueError("Unknown method. Use method='power_series' or method='frobenius'.")
