# ============================================================
# File: _solve_ode_autonomous_differential_equations.py
# Topic: Autonomous Differential Equations
#
# Autonomous means "no explicit t" (independent variable) in RHS:
#
# (A) First-order autonomous:
#       dy/dt = f(y)
#     Separable:
#       ∫ dy/f(y) = t + C
#     Optional IC: y(t0)=y0  -> solve for C
#
# (B) Second-order autonomous:
#       d2y/dt2 = F(y)
#     Reduce order by v = dy/dt (t eliminated):
#       y'' = dv/dt = (dv/dy)(dy/dt) = v dv/dy
#     So:
#       v dv/dy = F(y)
#       v dv = F(y) dy
#       (1/2) v^2 = ∫ F(y) dy + C1
#     Hence:
#       dy/dt = ± sqrt( 2(∫F(y) dy + C1) )
#     Then separable again:
#       ∫ dy / sqrt(2(∫F(y) dy + C1)) = ± t + C2
#
# Optional ICs for 2nd order:
#   - y(t0)=y0
#   - v(t0)=y'(t0)=v0  -> determines C1
#
# Notes:
#   - If SymPy can’t integrate, it keeps Integral(...)
#   - We return implicit solutions as much as possible.
# ============================================================

from __future__ import annotations

from typing import Optional, Union, Dict, Any

import sympy as sp


def solve_ode_autonomous_differential_equations(
    rhs: Union[str, sp.Expr],
    *,
    order: int = 1,
    t_symbol: str = "t",
    y_symbol: str = "y",
    C_symbol: str = "C",
    C1_symbol: str = "C_1",
    C2_symbol: str = "C_2",
    t0: Optional[Union[int, float, str, sp.Expr]] = None,
    y0: Optional[Union[int, float, str, sp.Expr]] = None,
    v0: Optional[Union[int, float, str, sp.Expr]] = None,  # for 2nd order: y'(t0)=v0
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Solve autonomous ODEs.

    Parameters
    ----------
    rhs:
        For order=1: f(y) in dy/dt = f(y)
        For order=2: F(y) in d2y/dt2 = F(y)
    order:
        1 or 2
    ICs:
        - order=1: (t0,y0)
        - order=2: (t0,y0) optional; (t0,y0,v0) helps determine C1

    Returns
    -------
    { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
    """
    if order not in (1, 2):
        raise ValueError("order must be 1 or 2 for this solver.")

    # -----------------------------
    # 1) Symbols + parsing
    # -----------------------------
    t = sp.Symbol(t_symbol, real=True)
    y = sp.Function(y_symbol)
    yt = y(t)

    yy = sp.Symbol(y_symbol, real=True)  # algebraic y for integration
    v = sp.Symbol("v", real=True)

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

    base_locals = {
        t_symbol: t,
        y_symbol: yy,
        "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 L(expr: sp.Expr) -> str:
        return sp.latex(expr)

    try:
        rhs_raw = sp.sympify(rhs, locals=base_locals) if isinstance(rhs, str) else sp.sympify(rhs)
    except Exception as e:
        raise ValueError(f"Could not parse rhs={rhs!r}. Error: {e}") from e

    F = sp.simplify(rhs_raw) if simplify else rhs_raw  # f(y) or F(y)

    use_ic_1 = (t0 is not None) and (y0 is not None)
    use_ic_2_energy = (t0 is not None) and (y0 is not None) and (v0 is not None)

    parsed_t0 = parsed_y0 = parsed_v0 = None
    if t0 is not None:
        parsed_t0 = sp.sympify(t0, locals=base_locals)
    if y0 is not None:
        parsed_y0 = sp.sympify(y0, locals=base_locals)
    if v0 is not None:
        parsed_v0 = sp.sympify(v0, locals=base_locals)

    # -----------------------------
    # 2) Build problem latex
    # -----------------------------
    if order == 1:
        ode_ltx = f"\\frac{{d{L(yt)}}}{{d{L(t)}}} = {L(F.subs(yy, yt))}"
        problem_s = f"Solve \\( {ode_ltx} \\)" + (
            f" given \\( {L(yt)}\\big|_{{{L(t)}={L(parsed_t0)}}}={L(parsed_y0)} \\)"
            if use_ic_1 else ""
        )

        # -----------------------------
        # 3) Solve: ∫ dy/f(y) = t + C
        # -----------------------------
        integrand = sp.simplify(1 / F) if simplify else (1 / F)
        I_y = sp.integrate(integrand, yy)  # ∫ dy/f(y)
        if simplify:
            I_y = sp.simplify(I_y)

        implicit_general = sp.Eq(I_y.subs(yy, yt), t + C)

        # IC: I(y0) = t0 + C => C = I(y0) - t0
        C_value = None
        implicit_ic = None
        if use_ic_1:
            try:
                Iy0 = I_y.subs(yy, parsed_y0)
                C_value_raw = Iy0 - parsed_t0
                C_value = sp.simplify(C_value_raw) if simplify else C_value_raw
                implicit_ic = sp.Eq(I_y.subs(yy, yt), t + C_value)
            except Exception:
                C_value = None
                implicit_ic = None

        # -----------------------------
        # 4) HTML solution
        # -----------------------------
        notes = []
        if bool(I_y.has(sp.Integral)):
            notes.append("SymPy kept an integral unevaluated in \\(\\int \\frac{dy}{f(y)}\\).")
        if F.free_symbols - {yy}:
            notes.append("Note: RHS should depend only on \\(y\\) for a strictly autonomous scalar ODE.")
        note_s = " ".join(notes)

        step1 = f"\\frac{{dy}}{{dt}}=f(y)\\Rightarrow \\frac{{dy}}{{f(y)}}=dt"
        step2 = f"\\int \\frac{{dy}}{{f(y)}}=\\int dt"
        step3 = f"{L(I_y)} = {L(t)} + {L(C)}"
        final_eq = implicit_ic if (implicit_ic is not None) else implicit_general

        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Autonomous ODE (First Order): \\(y' = f(y)\\)<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            f"& \\displaystyle {ode_ltx} \\\\[8pt]"
            f"& \\displaystyle \\Rightarrow {step1} \\\\"
            f"& \\displaystyle \\Rightarrow {step2} \\\\"
            f"& \\displaystyle \\Rightarrow {step3} \\\\[8pt]"
            f"& \\displaystyle \\Rightarrow \\text{{Implicit solution: }} {L(final_eq.lhs)} = {L(final_eq.rhs)} \\\\"
            "\\end{aligned}\\)"
            f"<div class='note'>{note_s}</div>"
            "</div></div>"
        )

        answer_s = f"\\( {L(final_eq.lhs)} = {L(final_eq.rhs)} \\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # =============================
    # order == 2
    # =============================
    ode_ltx = f"\\frac{{d^2{L(yt)}}}{{d{L(t)}^2}} = {L(F.subs(yy, yt))}"
    problem_s = f"Solve \\( {ode_ltx} \\)" + (
        (
            f" given \\( {L(yt)}\\big|_{{{L(t)}={L(parsed_t0)}}}={L(parsed_y0)} \\)"
            + (
                f" and \\( {L(sp.diff(yt, t))}\\big|_{{{L(t)}={L(parsed_t0)}}}={L(parsed_v0)} \\)"
                if use_ic_2_energy else ""
            )
        )
        if (t0 is not None and y0 is not None) else ""
    )

    # -----------------------------
    # Energy integral:
    # (1/2) v^2 = ∫F(y) dy + C1
    # -----------------------------
    W = sp.integrate(F, yy)  # ∫F(y) dy
    if simplify:
        W = sp.simplify(W)

    energy_eq = sp.Eq(sp.Rational(1, 2) * v**2, W + C1)

    # Determine C1 if v0 and y0 given
    C1_value = None
    if use_ic_2_energy:
        try:
            W_y0 = W.subs(yy, parsed_y0)
            C1_value_raw = sp.Rational(1, 2) * (parsed_v0**2) - W_y0
            C1_value = sp.simplify(C1_value_raw) if simplify else C1_value_raw
        except Exception:
            C1_value = None

    # dy/dt = ± sqrt(2(W + C1))
    radicand = 2 * (W + (C1_value if C1_value is not None else C1))
    if simplify:
        radicand = sp.simplify(radicand)

    v_expr = sp.sqrt(radicand)  # principal; we'll keep ± in text
    sep_integrand = sp.simplify(1 / v_expr) if simplify else (1 / v_expr)
    J_y = sp.integrate(sep_integrand, yy)
    if simplify:
        J_y = sp.simplify(J_y)

    # Implicit: ∫ dy / sqrt(2(W + C1)) = ± t + C2
    implicit_general = sp.Eq(J_y.subs(yy, yt), sp.Symbol("\\pm") * t + C2)

    # If we also have y(t0)=y0, we can fold into C2 (best-effort), but ± complicates.
    # We'll still compute a “+” branch constant for display if possible.
    C2_value_plus = None
    implicit_plus_ic = None
    if (t0 is not None) and (y0 is not None):
        try:
            Jy0 = J_y.subs(yy, parsed_y0)
            C2_value_plus_raw = Jy0 - parsed_t0
            C2_value_plus = sp.simplify(C2_value_plus_raw) if simplify else C2_value_plus_raw
            implicit_plus_ic = sp.Eq(J_y.subs(yy, yt), t + C2_value_plus)
        except Exception:
            C2_value_plus = None
            implicit_plus_ic = None

    # -----------------------------
    # HTML solution
    # -----------------------------
    notes = []
    if bool(W.has(sp.Integral)):
        notes.append("SymPy kept \\(\\int F(y)\\,dy\\) unevaluated.")
    if bool(J_y.has(sp.Integral)):
        notes.append("SymPy kept the final separation integral unevaluated.")
    note_s = " ".join(notes)

    step_r = (
        "\\text{Let } v=\\frac{dy}{dt}.\\;"
        "y''=\\frac{dv}{dt}=\\frac{dv}{dy}\\frac{dy}{dt}=v\\frac{dv}{dy}"
    )
    step_e1 = f"v\\frac{{dv}}{{dy}} = {L(F)}"
    step_e2 = f"v\\,dv = {L(F)}\\,dy"
    step_e3 = f"\\frac12 v^2 = \\int {L(F)}\\,dy + C_1"
    step_e4 = f"\\frac12 v^2 = {L(W)} + {L(C1_value if C1_value is not None else C1)}"
    step_sep1 = f"\\frac{{dy}}{{dt}} = \\pm \\sqrt{{2\\left({L(W)} + {L(C1_value if C1_value is not None else C1)}\\right)}}"
    step_sep2 = (
        f"\\Rightarrow \\int \\frac{{dy}}{{\\sqrt{{2\\left({L(W)} + {L(C1_value if C1_value is not None else C1)}\\right)}}}}"
        f" = \\pm t + C_2"
    )

    final_for_answer = implicit_plus_ic if (implicit_plus_ic is not None) else implicit_general

    # IMPORTANT FIX:
    # build optional IC line separately, with proper escaped braces in \\text{{...}}
    ic_line = ""
    if implicit_plus_ic is not None:
        ic_line = (
            f"& \\displaystyle \\Rightarrow \\text{{(Using }}{L(yt)}\\left({L(parsed_t0)}\\right)={L(parsed_y0)}"
            f"\\text{{, + branch): }} {L(final_for_answer.lhs)} = {L(final_for_answer.rhs)} \\\\"
        )

    soln_s = (
        "<div class='wrap'><div class='card'>"
        "<div class='badge'>Autonomous ODE (Second Order): \\(y'' = F(y)\\) (order reduction)<br/></div>"
        "\\(\\displaystyle \\begin{aligned}"
        f"& \\displaystyle {ode_ltx} \\\\[8pt]"
        f"& \\displaystyle \\Rightarrow {step_r} \\\\[6pt]"
        f"& \\displaystyle \\Rightarrow {step_e1} \\\\"
        f"& \\displaystyle \\Rightarrow {step_e2} \\\\"
        f"& \\displaystyle \\Rightarrow {step_e3} \\\\"
        f"& \\displaystyle \\Rightarrow {step_e4} \\\\[10pt]"
        f"& \\displaystyle \\Rightarrow {step_sep1} \\\\"
        f"& \\displaystyle {step_sep2} \\\\[8pt]"
        f"& \\displaystyle \\Rightarrow \\text{{Implicit solution (general): }} {L(J_y.subs(yy, yt))} = \\pm {L(t)} + {L(C2)} \\\\"
        f"{ic_line}"
        "\\end{aligned}\\)"
        f"<div class='note'>{note_s}</div>"
        "</div></div>"
    )

    answer_s = f"\\( {L(final_for_answer.lhs)} = {L(final_for_answer.rhs)} \\)"
    return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}
