# ============================================================
# File: _solve_ode_first_order_riccati.py
# Topic: First-Order Riccati ODE
#
# Form:
#   y' = q0(x) + q1(x)*y + q2(x)*y^2     (q2 != 0)
#
# Standard method (given a particular solution y_p(x)):
#   Substitute:  y = y_p + 1/v
#   Then v satisfies a linear ODE:
#     v' + (q1 + 2*q2*y_p)*v + q2 = 0
#
# Then:
#   Solve for v(x), and recover y(x) = y_p(x) + 1/v(x).
#
# Optional IC: y(x0)=y0, solve for constant.
#
# Notes:
#   - This solver expects a usable particular solution y_p(x).
#   - If SymPy can’t integrate, it keeps Integral(...) in the result.
# ============================================================

from __future__ import annotations

from typing import Any, Dict, Optional, Union

import sympy as sp


def solve_ode_first_order_riccati(
    q0_of_x: Union[str, sp.Expr],
    q1_of_x: Union[str, sp.Expr],
    q2_of_x: Union[str, sp.Expr],
    *,
    y_particular: Union[str, sp.Expr],
    x_symbol: str = "x",
    y_symbol: str = "y",
    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 Riccati:
        y' = q0(x) + q1(x)*y + q2(x)*y^2
    given a particular solution y_p(x).

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

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

    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 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)

    q0_raw = _parse(q0_of_x)
    q1_raw = _parse(q1_of_x)
    q2_raw = _parse(q2_of_x)
    yp_raw = _parse(y_particular)

    q0 = sp.simplify(q0_raw) if simplify else q0_raw
    q1 = sp.simplify(q1_raw) if simplify else q1_raw
    q2 = sp.simplify(q2_raw) if simplify else q2_raw
    yp = sp.simplify(yp_raw) if simplify else yp_raw

    # Guard: q2 != 0
    if simplify:
        if sp.simplify(q2) == 0:
            raise ValueError("Riccati requires q2(x) != 0 (quadratic term).")
    else:
        if q2 == 0:
            raise ValueError("Riccati requires q2(x) != 0 (quadratic term).")

    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) Validate particular solution (best-effort)
    # -----------------------------
    yp_prime = sp.diff(yp, x)
    rhs_at_yp = sp.simplify(q0 + q1 * yp + q2 * yp**2) if simplify else (q0 + q1 * yp + q2 * yp**2)
    residual = sp.simplify(yp_prime - rhs_at_yp) if simplify else (yp_prime - rhs_at_yp)

    # -----------------------------
    # 3) Substitute y = yp + 1/v  -> linear ODE for v
    # -----------------------------
    v = sp.Function("v")
    vx = v(x)

    P = sp.simplify(q1 + 2 * q2 * yp) if simplify else (q1 + 2 * q2 * yp)
    Q = sp.simplify(q2) if simplify else q2

    # v' + P v + Q = 0  => v' + P v = -Q
    mu = sp.exp(sp.integrate(P, x))
    I = sp.integrate(-Q * mu, x)
    v_general = sp.simplify((I + C) / mu) if simplify else (I + C) / mu

    # y = yp + 1/v
    y_general = sp.simplify(yp + 1 / v_general) if simplify else (yp + 1 / v_general)

    # -----------------------------
    # 4) Apply IC if available
    # -----------------------------
    C_value = None
    y_with_C = None
    if use_ic:
        try:
            # y(x0) = y0 => solve for C
            eqC = sp.Eq(y_general.subs(x, parsed_x0), parsed_y0)
            solC = sp.solve(eqC, C)
            if solC:
                C_value = sp.simplify(solC[0]) if simplify else solC[0]
                v_ic = sp.simplify((I + C_value) / mu) if simplify else (I + C_value) / mu
                y_with_C = sp.simplify(yp + 1 / v_ic) if simplify else (yp + 1 / v_ic)
        except Exception:
            C_value = None
            y_with_C = None

    # -----------------------------
    # 5) Build HTML (match your style)
    # -----------------------------
    ode_ltx = (
        f"\\frac{{d {L(yx)}}}{{d {L(x)}}}"
        f"= {L(q0)} + \\left({L(q1)}\\right){L(yx)} + \\left({L(q2)}\\right){L(yx)}^2"
    )

    # Substitution steps
    sub_ltx = f"{L(yx)} = {L(yp)} + \\frac{{1}}{{v({L(x)})}}"
    v_ode_ltx = f"v' + \\left({L(P)}\\right)v + {L(Q)} = 0"
    mu_ltx = f"\\mu(x)=\\exp\\left(\\int {L(P)}\\,d{L(x)}\\right) = {L(mu)}"
    v_sol_ltx = f"v(x)=\\frac{{\\int \\left(-{L(Q)}\\right)\\mu(x)\\,d{L(x)} + {L(C)}}}{{\\mu(x)}} = {L(v_general)}"

    if y_with_C is not None and C_value is not None:
        y_ans_ltx = f"{L(yx)} = {L(y_with_C)}"
        ic_ltx = f"{L(y)}({L(parsed_x0)})={L(parsed_y0)}"
    else:
        y_ans_ltx = f"{L(yx)} = {L(y_general)}"
        ic_ltx = None

    # Notes
    note_parts = []
    if residual != 0:
        note_parts.append(
            "Given \\(y_p\\) does not simplify to an exact particular solution under current simplification; proceeding symbolically."
        )
    if mu.has(sp.Integral):
        note_parts.append("SymPy kept an integral inside \\(\\mu\\) unevaluated.")
    if v_general.has(sp.Integral) or y_general.has(sp.Integral):
        note_parts.append("SymPy kept one or more integrals unevaluated in the final expression.")
    anti_note = " ".join(note_parts) if note_parts else ""

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

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

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

        & \\displaystyle \\Rightarrow \\text{{Given a particular solution }} y_p(x)={L(yp)} \\\\
        & \\displaystyle \\Rightarrow \\text{{Use substitution:}}\\; {sub_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow \\text{{This gives a linear ODE for }}v: \\\\
        & \\displaystyle \\Rightarrow {v_ode_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Integrating factor:}} \\\\
        & \\displaystyle \\Rightarrow {mu_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow \\text{{Solve for }}v(x): \\\\
        & \\displaystyle \\Rightarrow {v_sol_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Back-substitute:}}\\; {y_ans_ltx} \\\\

        \\end{{aligned}}
        \\)
        <div class="note">{anti_note}</div>

      </div>
    </div>
    """.replace("\n", "")

    # IC block (if solved)
    if ic_ltx is not None and C_value is not None:
        ic_block = f"""
        <br/>
        \\(\\displaystyle
        \\begin{{aligned}}
        & \\displaystyle \\text{{Apply IC }} {ic_ltx}: \\\\
        & \\displaystyle \\Rightarrow {L(C)} = {L(C_value)} \\\\
        & \\displaystyle \\Rightarrow {L(yx)} = {L(y_with_C)} \\\\
        \\end{{aligned}}
        \\)
        """.replace("\n", "")
        soln_s = soln_stmt + ic_block
        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {ic_ltx} \\)"
        answer_s = f"\\( {L(yx)} = {L(y_with_C)} \\)"
    else:
        soln_s = soln_stmt
        problem_s = f"Solve \\( {ode_ltx} \\)"
        answer_s = f"\\( {y_ans_ltx} \\)"

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



