# ============================================================
# File: _solve_ode_first_order_lagrange.py
# Topic: First-Order Lagrange ODE
#
# Form:
#   y = x*phi(p) + psi(p),    where p = dy/dx
#
# Differentiate:
#   y' = p = phi(p) + (x*phi'(p) + psi'(p)) * dp/dx
# => (x*phi'(p) + psi'(p)) * dp/dx = p - phi(p)
#
# Invert:
#   dx/dp = (x*phi'(p) + psi'(p)) / (p - phi(p))
# => dx/dp - [phi'(p)/(p - phi(p))] x = psi'(p)/(p - phi(p))
#
# This is a linear ODE in x(p):
#   x'(p) + A(p) x = B(p),
#   where A(p) = -phi'(p)/(p - phi(p)),  B(p) = psi'(p)/(p - phi(p))
#
# Then parametric solution:
#   x(p) = ( ∫ μ(p) B(p) dp + C ) / μ(p)
#   y(p) = x(p)*phi(p) + psi(p)
#
# Additional straight-line solutions:
#   If p = phi(p) has roots p=c, then y = c x + psi(c)
#
# Optional IC (best-effort):
#   y(x0)=y0:
#     solve y0 = x0*phi(p0)+psi(p0) for p0,
#     then use x(p0)=x0 to solve for C.
# ============================================================

from __future__ import annotations

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

import sympy as sp


def solve_ode_first_order_lagrange(
    phi_of_p: Union[str, sp.Expr],
    psi_of_p: Union[str, sp.Expr],
    *,
    x_symbol: str = "x",
    y_symbol: str = "y",
    p_symbol: str = "p",
    C_symbol: str = "C",
    x0: Optional[Union[int, float, str, sp.Expr]] = None,
    y0: Optional[Union[int, float, str, sp.Expr]] = None,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Solve Lagrange ODE:
        y = x*phi(p) + psi(p),  p = dy/dx

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

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

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

    C = sp.Symbol(C_symbol, real=True)

    base_locals = {
        x_symbol: x,
        p_symbol: p,
        "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)

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

    phi_raw = _parse(phi_of_p)
    psi_raw = _parse(psi_of_p)

    phi = sp.simplify(phi_raw) if simplify else phi_raw
    psi = sp.simplify(psi_raw) if simplify else psi_raw

    use_ic = (x0 is not None) and (y0 is not None)
    parsed_x0 = parsed_y0 = None
    if use_ic:
        parsed_x0 = _parse(x0)
        parsed_y0 = _parse(y0)

    # -----------------------------
    # 2) Build linear ODE for x(p)
    # -----------------------------
    denom = sp.simplify(p - phi) if simplify else (p - phi)
    phi_p = sp.diff(phi, p)
    psi_p = sp.diff(psi, p)

    # A(p) = -phi'(p)/(p-phi(p)),  B(p) = psi'(p)/(p-phi(p))
    A = sp.simplify(-phi_p / denom) if simplify else (-phi_p / denom)
    B = sp.simplify(psi_p / denom) if simplify else (psi_p / denom)

    # Integrating factor μ(p) = exp(∫A dp)
    mu = sp.exp(sp.integrate(A, p))
    I = sp.integrate(mu * B, p)

    x_general = sp.simplify((I + C) / mu) if simplify else (I + C) / mu
    y_general = sp.simplify(x_general * phi + psi) if simplify else (x_general * phi + psi)

    # -----------------------------
    # 3) Straight-line solutions: p = phi(p)
    # -----------------------------
    line_solutions: List[sp.Expr] = []
    try:
        roots = sp.solve(sp.Eq(p, phi), p)
        for r in roots:
            rr = sp.simplify(r) if simplify else r
            # y = r x + psi(r)
            line_solutions.append(sp.simplify(rr * x + psi.subs(p, rr)) if simplify else (rr * x + psi.subs(p, rr)))
    except Exception:
        pass

    # -----------------------------
    # 4) Apply IC (best-effort)
    # -----------------------------
    p0 = None
    C_value = None
    x_ic = None
    y_ic = None

    if use_ic:
        # First, solve y0 = x0*phi(p0) + psi(p0) for p0 (no C here)
        try:
            eq_p0 = sp.Eq(parsed_y0, parsed_x0 * phi + psi)
            sol_p0 = sp.solve(eq_p0, p)
            if sol_p0:
                p0 = sp.simplify(sol_p0[0]) if simplify else sol_p0[0]

                # Then use x(p0)=x0 to solve for C:
                # x(p) = (I(p)+C)/mu(p) => C = x0*mu(p0) - I(p0)
                C_value_raw = parsed_x0 * mu.subs(p, p0) - I.subs(p, p0)
                C_value = sp.simplify(C_value_raw) if simplify else C_value_raw

                x_ic = sp.simplify(x_general.subs(C, C_value)) if simplify else x_general.subs(C, C_value)
                y_ic = sp.simplify(y_general.subs(C, C_value)) if simplify else y_general.subs(C, C_value)
        except Exception:
            p0 = None
            C_value = None

    # -----------------------------
    # 5) HTML formatting (match your style)
    # -----------------------------
    ode_ltx = f"{L(yx)} = {L(x)}\\,\\phi(p) + \\psi(p),\\quad p=\\frac{{d{L(yx)}}}{{d{L(x)}}}"
    phi_ltx = f"\\phi(p)={L(phi)}"
    psi_ltx = f"\\psi(p)={L(psi)}"

    diff_ltx = (
        f"p = \\phi(p) + \\big(x\\phi'(p)+\\psi'(p)\\big)\\,\\frac{{dp}}{{d{L(x)}}}"
    )
    inv_ltx = (
        f"\\frac{{d{L(x)}}}{{dp}} = \\frac{{x\\phi'(p)+\\psi'(p)}}{{p-\\phi(p)}}"
    )
    lin_ltx = (
        f"\\frac{{d{L(x)}}}{{dp}} - \\frac{{\\phi'(p)}}{{p-\\phi(p)}}\\,{L(x)}"
        f" = \\frac{{\\psi'(p)}}{{p-\\phi(p)}}"
    )
    A_ltx = f"A(p)={L(A)}"
    B_ltx = f"B(p)={L(B)}"
    mu_ltx = f"\\mu(p)=\\exp\\left(\\int A(p)\\,dp\\right) = {L(mu)}"
    x_sol_ltx = f"x(p)=\\frac{{\\int \\mu(p)B(p)\\,dp + {L(C)}}}{{\\mu(p)}} = {L(x_general)}"
    y_sol_ltx = f"y(p)=x(p)\\,\\phi(p)+\\psi(p) = {L(y_general)}"

    note_parts = []
    if mu.has(sp.Integral):
        note_parts.append("SymPy kept an integral inside \\(\\mu(p)\\) unevaluated.")
    if x_general.has(sp.Integral) or y_general.has(sp.Integral):
        note_parts.append("SymPy kept one or more integrals unevaluated in the parametric solution.")
    if line_solutions:
        note_parts.append("Also includes straight-line solutions from roots of \\(p=\\phi(p)\\).")
    anti_note = " ".join(note_parts)

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: First-Order Lagrange ODE<br/></div>

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

        & \\displaystyle {ode_ltx} \\\\[6pt]
        & \\displaystyle {phi_ltx},\\; {psi_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Differentiate w.r.t. }}x: \\\\
        & \\displaystyle \\Rightarrow {diff_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow \\big(x\\phi'(p)+\\psi'(p)\\big)\\,\\frac{{dp}}{{dx}} = p-\\phi(p) \\\\
        & \\displaystyle \\Rightarrow {inv_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow \\text{{Linear ODE in }}x(p): \\\\
        & \\displaystyle \\Rightarrow {lin_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow {A_ltx},\\quad {B_ltx} \\\\
        & \\displaystyle \\Rightarrow {mu_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow {x_sol_ltx} \\\\
        & \\displaystyle \\Rightarrow {y_sol_ltx} \\\\

        \\end{{aligned}}
        \\)

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

    # Add straight-line solutions block
    extra_block = ""
    if line_solutions:
        lines_ltx = " \\\\ ".join([f"& \\displaystyle y = {L(expr)}" for expr in line_solutions])
        extra_block = f"""
        <br/>
        \\(\\displaystyle
        \\begin{{aligned}}
        & \\displaystyle \\text{{Straight-line solutions from }} p=\\phi(p): \\\\
        {lines_ltx}
        \\end{{aligned}}
        \\)
        """.replace("\n", "")

    # IC block (if we managed to solve p0 and C)
    if use_ic and (p0 is not None) and (C_value is not None) and (x_ic is not None) and (y_ic is not None):
        ic_ltx = f"{L(y)}({L(parsed_x0)})={L(parsed_y0)}"
        ic_block = f"""
        <br/>
        \\(\\displaystyle
        \\begin{{aligned}}
        & \\displaystyle \\text{{Apply IC }} {ic_ltx}: \\\\
        & \\displaystyle \\Rightarrow p_0 = {L(p0)} \\\\
        & \\displaystyle \\Rightarrow {L(C)} = {L(C_value)} \\\\
        & \\displaystyle \\Rightarrow x(p) = {L(x_ic)} \\\\
        & \\displaystyle \\Rightarrow y(p) = {L(y_ic)} \\\\
        \\end{{aligned}}
        \\)
        """.replace("\n", "")
        soln_s = soln_stmt + extra_block + ic_block
        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {ic_ltx} \\)"
        answer_s = f"\\( x(p)={L(x_ic)},\\; y(p)={L(y_ic)} \\)"
    else:
        soln_s = soln_stmt + extra_block
        problem_s = f"Solve \\( {ode_ltx} \\)"
        answer_s = f"\\( x(p)={L(x_general)},\\; y(p)={L(y_general)} \\)"

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


