# ============================================================
# File: _solve_ode_conservative_newtonian_mechanics_form.py
# Topic: Conservative / Newtonian Mechanics Form (y'' = f(y))
# Typical conservative form: y'' = -V'(y)
# Method (energy integral):
#   Multiply by y':
#       y'' y' = f(y) y'
#   Then:
#       d/dx ( 1/2 (y')^2 ) = d/dx ( ∫ f(y) dy )
#   So:
#       1/2 (y')^2 = ∫ f(y) dy + C
#   If we define potential V(y) = -∫ f(y) dy, then:
#       1/2 (y')^2 + V(y) = E
#   Implicit quadrature:
#       ∫ dy / sqrt( 2(E - V(y)) ) = ±(x - x0)   (mass=1)
#
# Optional IC: y(x0)=y0, y'(x0)=dy0
# ============================================================

from __future__ import annotations

from typing import Optional, Union, Dict, Any

import sympy as sp


def solve_ode_conservative_newtonian_mechanics_form(
    f_of_y: Union[str, sp.Expr],
    *,
    x_symbol: str = "x",
    y_symbol: str = "y",
    C_symbol: str = "C",
    E_symbol: str = "E",
    x0: Optional[Union[int, float, str, sp.Expr]] = None,
    y0: Optional[Union[int, float, str, sp.Expr]] = None,
    dy0: Optional[Union[int, float, str, sp.Expr]] = None,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Solve the autonomous second-order ODE:
        y'' = f(y)
    with the conservative/energy method.

    Returns HTML steps with:
      - first integral: 1/2 (y')^2 = ∫ f(y) dy + C
      - potential form: 1/2 (y')^2 + V(y) = E where V(y) = -∫ f(y) dy
      - implicit quadrature: ∫ dy/sqrt(2(E - V(y))) = ±(x - x0)

    Optional IC:
        y(x0)=y0, y'(x0)=dy0
      gives E (or C).

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

    # -----------------------------
    # 1) Symbols + parsing
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    y = sp.Function(y_symbol)
    yx = y(x)

    yy = sp.Symbol(y_symbol, real=True)  # for integrating in y cleanly
    C = sp.Symbol(C_symbol)
    E = sp.Symbol(E_symbol)

    base_locals = {
        x_symbol: x,
        y_symbol: yy,  # parse f(y) as expression in 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.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,
    }
    if locals_map:
        base_locals.update(locals_map)

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

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

    f_y = sp.simplify(f_raw) if simplify else f_raw

    # -----------------------------
    # 2) Energy integral
    # -----------------------------
    # F(y) = ∫ f(y) dy
    F_raw = sp.integrate(f_y, yy)
    F = sp.simplify(F_raw) if simplify else F_raw

    # potential V(y) = -F(y)
    V = sp.simplify(-F) if simplify else (-F)

    # First integral:
    # 1/2 (y')^2 = F(y) + C
    # Energy form:
    # 1/2 (y')^2 + V(y) = E, where E is a constant
    #
    # Relate constants: E = -C  (since V=-F) up to renaming.
    # We'll present both.

    # Optional IC
    use_ic = (x0 is not None) and (y0 is not None) and (dy0 is not None)
    parsed_x0 = parsed_y0 = parsed_dy0 = None
    E_value = None
    C_value = None

    if use_ic:
        parsed_x0 = sp.sympify(x0, locals=base_locals)
        parsed_y0 = sp.sympify(y0, locals=base_locals)
        parsed_dy0 = sp.sympify(dy0, locals=base_locals)

        # E = 1/2 (y')^2 + V(y)
        E_value_raw = sp.Rational(1, 2) * (parsed_dy0**2) + V.subs(yy, parsed_y0)
        E_value = sp.simplify(E_value_raw) if simplify else E_value_raw

        # Also C from 1/2 (y')^2 = F(y) + C
        C_value_raw = sp.Rational(1, 2) * (parsed_dy0**2) - F.subs(yy, parsed_y0)
        C_value = sp.simplify(C_value_raw) if simplify else C_value_raw

    # Implicit quadrature (mass=1):
    # dx = dy / y' , and from energy: y' = ±sqrt(2(E - V(y)))
    # => ∫ dy / sqrt(2(E - V(y))) = ±x + const
    denom = sp.sqrt(2 * (E - V))
    quad_integrand = sp.simplify(1 / denom) if simplify else (1 / denom)
    quad_raw = sp.integrate(quad_integrand, yy)
    quad = sp.simplify(quad_raw) if simplify else quad_raw

    note_parts = []
    if bool(F.has(sp.Integral)):
        note_parts.append("SymPy kept \\(\\int f(y)\\,dy\\) unevaluated.")
    if bool(quad.has(sp.Integral)):
        note_parts.append("SymPy kept the quadrature integral unevaluated.")
    anti_note = " ".join(note_parts)

    # -----------------------------
    # 3) Build HTML steps
    # -----------------------------
    yprime = sp.diff(yx, x)
    y2 = sp.diff(yx, x, 2)

    ode_ltx = f"{L(y2)} = {L(f_y.subs(yy, yx))}"

    mult_ltx = f"{L(y2)}\\,{L(yprime)} = {L(f_y.subs(yy, yx))}\\,{L(yprime)}"
    lhs_ltx = f"\\frac{{d}}{{d{L(x)}}}\\left(\\frac{{1}}{{2}}\\left({L(yprime)}\\right)^2\\right)"
    rhs_ltx = f"\\frac{{d}}{{d{L(x)}}}\\left(\\int {L(f_y.subs(yy, yx))}\\, d{L(yx)}\\right)"

    first_int_ltx = f"\\frac{{1}}{{2}}\\left({L(yprime)}\\right)^2 = \\int {L(f_y.subs(yy, yx))}\\, d{L(yx)} + {L(C)}"
    F_ltx = f"\\int {L(f_y)}\\, d{L(yy)} = {L(F)}"
    first_int_simpl_ltx = f"\\frac{{1}}{{2}}\\left({L(yprime)}\\right)^2 = {L(F.subs(yy, yx))} + {L(C)} \\qquad \\cdots \\text{{(i)}}"

    V_ltx = f"V({L(yy)}) = -\\int {L(f_y)}\\, d{L(yy)} = {L(V)}"
    energy_ltx = f"\\frac{{1}}{{2}}\\left({L(yprime)}\\right)^2 + V({L(yx)}) = {L(E)} \\qquad \\cdots \\text{{(ii)}}"

    quad_ltx = (
        f"\\int \\frac{{d{L(yx)}}}{{\\sqrt{{2\\left({L(E)} - V({L(yx)})\\right)}}}}"
        f" = \\pm {L(x)} + {L(C)}"
        f" \\qquad \\cdots \\text{{(Quadrature)}}"
    )

    soln_stmts = []
    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: Conservative / Newtonian Mechanics Form<br/></div>

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

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

        & \\displaystyle \\Rightarrow \\text{{Multiply both sides by }} {L(yprime)}: \\\\
        & \\displaystyle \\Rightarrow {mult_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow \\text{{Recognize total derivatives:}} \\\\
        & \\displaystyle \\Rightarrow {lhs_ltx} = {rhs_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow \\text{{Integrate w.r.t. }} {L(x)}: \\\\
        & \\displaystyle \\Rightarrow {first_int_ltx} \\\\
        & \\displaystyle \\Rightarrow {F_ltx} \\\\
        & \\displaystyle \\Rightarrow {first_int_simpl_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Define potential (so that }} y''=-V'(y)\\text{{):}} \\\\
        & \\displaystyle \\Rightarrow {V_ltx} \\\\
        & \\displaystyle \\Rightarrow {energy_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Hence an implicit solution via quadrature:}} \\\\
        & \\displaystyle \\Rightarrow {quad_ltx} \\\\

        \\end{{aligned}}
        \\)
        <div class="note">{anti_note}</div>
      </div>
    </div>
    """
    soln_stmts.append(soln_stmt.replace("\n", ""))

    # IC block (show constants; answer should NOT include "from (i)" style tags)
    if use_ic and (E_value is not None) and (C_value is not None):
        ic1_ltx = f"{L(yx)}\\Big|_{{{L(x)}={L(parsed_x0)}}} = {L(parsed_y0)}"
        ic2_ltx = f"{L(yprime)}\\Big|_{{{L(x)}={L(parsed_x0)}}} = {L(parsed_dy0)}"

        E_expr_ltx = f"{L(E)} = \\frac{{1}}{{2}}\\left({L(parsed_dy0)}\\right)^2 + V\\left({L(parsed_y0)}\\right)"
        E_val_ltx = f"{L(E)} = {L(E_value)} \\qquad \\text{{Given }} {ic1_ltx},\\; {ic2_ltx}"

        # Quadrature with E_value
        quad_ic_soln_ltx = (
            f"\\int \\frac{{d{L(yx)}}}{{\\sqrt{{2\\left({L(E_value)} - V({L(yx)})\\right)}}}}"
            f" = \\pm {L(x)} + {L(C)}"
        )

        ic_block = f"""
            <br/>
            \\(\\displaystyle
            \\begin{{aligned}}
            & \\displaystyle \\Rightarrow {E_expr_ltx} \\\\
            & \\displaystyle \\Rightarrow {E_val_ltx} \\\\[6pt]
            & \\displaystyle \\Rightarrow {quad_ic_soln_ltx} \\\\
            \\end{{aligned}}
            \\)
        """
        soln_stmts.append(ic_block.replace("\n", ""))

        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {ic1_ltx},\\; {ic2_ltx} \\)"
        # Final answer: energy equation with computed E (clean)
        answer_s = f"\\( \\frac{{1}}{{2}}\\left({L(yprime)}\\right)^2 + V({L(yx)}) = {L(E_value)} \\)"
    else:
        problem_s = f"Solve \\( {ode_ltx} \\)"
        answer_s = f"\\( {energy_ltx} \\)"

    soln_s = "".join(soln_stmts)

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

