# ============================================================
# File: _solve_ode_higher_order_constant_coeff.py
# Topic: Higher-Order Linear Homogeneous ODE (Constant Coefficients)
#
# Form:
#   a_n y^(n) + a_{n-1} y^(n-1) + ... + a_1 y' + a_0 y = 0
#
# Method:
#   1) Assume y = e^{rx}
#   2) Characteristic equation:
#        a_n r^n + a_{n-1} r^{n-1} + ... + a_0 = 0
#   3) Roots and multiplicities give the solution basis:
#      - Real root r with multiplicity m:
#          e^{rx}, x e^{rx}, ..., x^{m-1} e^{rx}
#      - Complex pair α ± iβ with multiplicity m:
#          e^{αx} x^k cos(βx),  e^{αx} x^k sin(βx),  k=0..m-1
#
# Notes:
#   - Best-effort root handling via sympy.roots
#   - Intended for factorable polynomials produced by generator
# ============================================================

from __future__ import annotations

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

import sympy as sp


def solve_ode_higher_order_constant_coeff(
    coeffs: List[Union[str, sp.Expr, int, float]],
    *,
    x_symbol: str = "x",
    y_symbol: str = "y",
    C_prefix: str = "C",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Solve: a_n y^(n) + ... + a_0 y = 0  (constant coefficients)

    coeffs: [a_n, a_{n-1}, ..., a_0]  (length n+1, a_n != 0)

    Returns:
        { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
    """

    # -----------------------------
    # 1) Symbols + parsing
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    r = sp.Symbol("r")

    y = sp.Function(y_symbol)
    yx = y(x)

    base_locals = {
        x_symbol: x,
        "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.acos,
        "atan": sp.atan,
        "sinh": sp.sinh,
        "cosh": sp.cosh,
        "tanh": sp.tanh,
        "Abs": sp.Abs,
        "sec": sp.sec,
        "csc": sp.csc,
        "cot": sp.cot,
        "Piecewise": sp.Piecewise,
        "Heaviside": sp.Heaviside,
    }
    if locals_map:
        base_locals.update(locals_map)

    def _parse(v):
        if isinstance(v, str):
            return sp.sympify(v, locals=base_locals)
        return sp.sympify(v)

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

    if not coeffs or len(coeffs) < 2:
        raise ValueError("coeffs must be a list [a_n, ..., a_0] with length >= 2.")

    a_list = [_parse(c) for c in coeffs]
    if simplify:
        a_list = [sp.simplify(a) for a in a_list]

    n = len(a_list) - 1
    a_n = a_list[0]
    if simplify:
        if sp.simplify(a_n) == 0:
            raise ValueError("Leading coefficient a_n must be nonzero.")
    else:
        if a_n == 0:
            raise ValueError("Leading coefficient a_n must be nonzero.")

    # -----------------------------
    # 2) Build ODE and characteristic polynomial
    # -----------------------------
    ode_expr = sp.Integer(0)
    for k, a_k in enumerate(a_list):
        order = n - k
        if order == 0:
            ode_expr += a_k * yx
        else:
            ode_expr += a_k * sp.diff(yx, x, order)

    # Characteristic polynomial: a_n r^n + ... + a_0
    char_poly = sp.Integer(0)
    for k, a_k in enumerate(a_list):
        power = n - k
        char_poly += a_k * r**power

    char_poly_s = sp.factor(char_poly) if simplify else char_poly

    # -----------------------------
    # 3) Roots with multiplicities (best-effort)
    # -----------------------------
    roots_dict = sp.roots(char_poly_s, r)  # {root: multiplicity}
    # If roots_dict is empty, fall back to Poly.all_roots (may return RootOf)
    if not roots_dict:
        try:
            poly = sp.Poly(char_poly_s, r)
            all_roots = poly.all_roots()
            # Convert list to multiplicity dict
            roots_dict = {}
            for rr in all_roots:
                roots_dict[rr] = roots_dict.get(rr, 0) + 1
        except Exception:
            roots_dict = {}

    if not roots_dict:
        # graceful fallback
        ode_ltx = _latex_ode_from_coeffs(a_list, x, yx, L)
        problem_s = f"Solve \\( {ode_ltx} \\) using the characteristic equation."
        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Note: Higher-Order Linear ODE (Constant Coefficients)<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            f"& \\displaystyle \\text{{Characteristic: }} {L(char_poly_s)}=0 \\\\"
            "& \\displaystyle \\Rightarrow \\text{Could not extract roots symbolically.} \\\\"
            "\\end{aligned}\\)"
            "</div></div>"
        )
        answer_s = "\\(\\text{Could not extract roots symbolically.}\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # -----------------------------
    # 4) Build basis solutions
    # -----------------------------
    processed: set = set()
    basis: List[sp.Expr] = []
    basis_labels: List[str] = []  # for explanation
    constants: List[sp.Symbol] = []

    def _is_real(z: sp.Expr) -> bool:
        # best-effort real check
        zz = sp.simplify(z) if simplify else z
        imz = sp.im(zz)
        if simplify:
            return sp.simplify(imz) == 0
        return imz == 0

    def _beta_pos(beta: sp.Expr) -> sp.Expr:
        b = sp.simplify(beta) if simplify else beta
        if b.is_negative is True:
            return -b
        # Keep it positive-ish
        return sp.Abs(b) if b.free_symbols else b

    # Sort roots for stable output (best-effort)
    roots_items = list(roots_dict.items())

    def _sort_key(item):
        rr, mm = item
        rr_s = sp.N(rr) if rr.is_number else rr
        return (str(rr_s), mm)

    roots_items.sort(key=_sort_key)

    # Build basis
    for root, mult in roots_items:
        if root in processed:
            continue

        if _is_real(root):
            # real root
            for k in range(mult):
                expr = sp.exp(root * x) * (x**k)
                expr = sp.simplify(expr) if simplify else expr
                basis.append(expr)
                basis_labels.append(f"real root {sp.latex(root)}, k={k}")
            processed.add(root)
        else:
            # complex root: handle conjugate pair once if present
            conj = sp.conjugate(root)
            if conj in roots_dict and conj not in processed:
                # multiplicities should match; use min to be safe
                m2 = roots_dict.get(conj, mult)
                m = min(mult, m2)
                alpha = sp.re(root)
                beta = _beta_pos(sp.im(root))

                for k in range(m):
                    epart = sp.exp(alpha * x) * (x**k)
                    cos_part = sp.simplify(epart * sp.cos(beta * x)) if simplify else (epart * sp.cos(beta * x))
                    sin_part = sp.simplify(epart * sp.sin(beta * x)) if simplify else (epart * sp.sin(beta * x))
                    basis.append(cos_part)
                    basis.append(sin_part)
                    basis_labels.append(f"complex pair {sp.latex(alpha)}±i{sp.latex(beta)}, cos, k={k}")
                    basis_labels.append(f"complex pair {sp.latex(alpha)}±i{sp.latex(beta)}, sin, k={k}")

                processed.add(root)
                processed.add(conj)
            else:
                # complex root without conjugate present (rare) — keep exp(root x) basis
                for k in range(mult):
                    expr = sp.exp(root * x) * (x**k)
                    expr = sp.simplify(expr) if simplify else expr
                    basis.append(expr)
                    basis_labels.append(f"complex root {sp.latex(root)}, k={k}")
                processed.add(root)

    # Constants C_1, C_2, ...
    for i in range(len(basis)):
        constants.append(sp.Symbol(f"{C_prefix}_{i+1}"))

    y_general = sp.Integer(0)
    for Ci, bi in zip(constants, basis):
        y_general += Ci * bi
    y_general = sp.simplify(y_general) if simplify else y_general

    # -----------------------------
    # 5) HTML formatting (match your style)
    # -----------------------------
    ode_ltx = _latex_ode_from_coeffs(a_list, x, yx, L)

    char_ltx = f"{L(char_poly_s)}=0"

    # Roots display
    roots_lines = []
    for rr, mm in roots_items:
        roots_lines.append(f"& \\displaystyle r={L(rr)} \\;\\; (\\text{{mult}}={mm})")
    roots_block = " \\\\ ".join(roots_lines)

    # Basis display (short)
    basis_lines = []
    for Ci, bi in zip(constants, basis):
        basis_lines.append(f"& \\displaystyle {L(Ci)}\\,{L(bi)}")
    basis_block = " \\\\ ".join(basis_lines[:14])  # keep it compact
    if len(basis_lines) > 14:
        basis_block += " \\\\ & \\displaystyle \\cdots"

    note_parts = []
    if any([char_poly_s.has(sp.RootOf), char_poly_s.has(sp.AlgebraicNumber)]):
        note_parts.append("Some roots may be represented implicitly.")
    anti_note = " ".join(note_parts)

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: Higher-Order Linear ODE (Constant Coefficients)<br/></div>

        \\(\\displaystyle
        \\begin{{aligned}}

        & \\displaystyle {ode_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Assume }} y=e^{{rx}} \\Rightarrow \\text{{Characteristic equation:}} \\\\
        & \\displaystyle \\Rightarrow {char_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Roots (with multiplicity):}} \\\\
        {roots_block} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{General solution from root structure:}} \\\\
        & \\displaystyle {L(yx)} = {L(y_general)} \\\\

        \\end{{aligned}}
        \\)

        <div class="note">{anti_note}</div>
      </div>
    </div>
    """.replace("\n", "")

    problem_s = f"Solve \\( {ode_ltx} \\) using the characteristic equation."
    answer_s = f"\\( {L(yx)} = {L(y_general)} \\)"

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


def _latex_ode_from_coeffs(a_list: List[sp.Expr], x: sp.Symbol, yx: sp.Expr, L) -> str:
    """
    Helper: Build latex for a_n y^(n) + ... + a_0 y = 0
    """
    n = len(a_list) - 1
    terms = []
    for k, a_k in enumerate(a_list):
        order = n - k
        if a_k == 0:
            continue
        if order == 0:
            t = f"\\left({L(a_k)}\\right){L(yx)}"
        elif order == 1:
            t = f"\\left({L(a_k)}\\right)\\frac{{d{L(yx)}}}{{d{L(x)}}}"
        else:
            t = f"\\left({L(a_k)}\\right)\\frac{{d^{order}{L(yx)}}}{{d{L(x)}^{order}}}"
        terms.append(t)

    if not terms:
        return "0=0"

    return " + ".join(terms) + " = 0"
