# ============================================================
# File: _solve_ode_first_order_bernoulli.py
# Topic: First-order Bernoulli ODE (full-steps HTML)
# Form:  y' + P(x) y = Q(x) y^n ,  n != 0,1
# Output: dict {problem_s, soln_s, answer_s}
# ============================================================

from __future__ import annotations

from typing import Optional, Union, Dict, Any

import sympy as sp


def solve_ode_first_order_bernoulli(
    P: Union[str, sp.Expr],
    Q: Union[str, sp.Expr],
    n: Union[int, float, 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:
    """
    Full-steps HTML solution for Bernoulli ODE:
        y' + P(x) y = Q(x) y^n,   n != 0,1

    Substitution:
        v = y^(1-n)
      => v' = (1-n) y^(-n) y'

    After multiplying the ODE by y^(-n), it becomes linear in v:
        v' + (1-n)P(x) v = (1-n)Q(x)

    Optional initial condition:
        y(x0) = y0

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

    # -----------------------------
    # 1) Setup symbols + parsing
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    yx = sp.Function(y_symbol)(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,
    }
    if locals_map:
        base_locals.update(locals_map)

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

    try:
        P_ = sp.sympify(P, locals=base_locals) if isinstance(P, str) else sp.sympify(P)
        Q_ = sp.sympify(Q, locals=base_locals) if isinstance(Q, str) else sp.sympify(Q)
        n_ = sp.sympify(n, locals=base_locals) if isinstance(n, str) else sp.sympify(n)
    except Exception as e:
        raise ValueError(f"Could not parse inputs: P={P!r}, Q={Q!r}, n={n!r}. Error: {e}") from e

    if simplify:
        P_, Q_, n_ = sp.simplify(P_), sp.simplify(Q_), sp.simplify(n_)

    # Guard: n must not be 0 or 1
    if (n_ - 1).equals(0) or (n_ - 0).equals(0):
        raise ValueError("Bernoulli ODE requires n != 0 and n != 1.")

    one_minus_n = sp.simplify(1 - n_) if simplify else (1 - n_)

    # -----------------------------
    # 2) Core math (Bernoulli -> linear)
    # -----------------------------
    # Define v = y^(1-n)
    v = sp.Function("v")(x)

    # Linear ODE in v:  v' + (1-n)P v = (1-n)Q
    P_lin = sp.simplify(one_minus_n * P_) if simplify else (one_minus_n * P_)
    Q_lin = sp.simplify(one_minus_n * Q_) if simplify else (one_minus_n * Q_)

    # Integrating factor for v' + P_lin v = Q_lin
    intP_lin_raw = sp.integrate(P_lin, x)
    intP_lin = sp.simplify(intP_lin_raw) if simplify else intP_lin_raw

    mu_raw = sp.exp(intP_lin)
    mu = sp.simplify(mu_raw) if simplify else mu_raw

    muQ_raw = sp.simplify(mu * Q_lin) if simplify else (mu * Q_lin)
    I_raw = sp.integrate(muQ_raw, x)
    I = sp.simplify(I_raw) if simplify else I_raw

    # v = (I + C)/mu
    v_general = (I + C) / mu
    if simplify:
        v_general = sp.simplify(v_general)

    # y = v^(1/(1-n))
    power_inv = sp.simplify(1 / one_minus_n) if simplify else (1 / one_minus_n)
    y_general = sp.Pow(v_general, power_inv)
    if simplify:
        y_general = sp.simplify(y_general)

    # IC (optional): y(x0)=y0 -> v(x0)=y0^(1-n)
    use_ic = (x0 is not None) and (y0 is not None)
    C_value: Optional[sp.Expr] = None
    y_particular: Optional[sp.Expr] = None

    parsed_x0 = None
    parsed_y0 = None
    v0 = None
    if use_ic:
        parsed_x0 = sp.sympify(x0, locals=base_locals)
        parsed_y0 = sp.sympify(y0, locals=base_locals)
        v0 = sp.simplify(sp.Pow(parsed_y0, one_minus_n)) if simplify else sp.Pow(parsed_y0, one_minus_n)

        # From mu*v = I + C  =>  C = mu*v - I
        C_value_raw = (mu * v_general - I).subs({x: parsed_x0, C: 0})  # v_general has +C; set C=0 gives I/mu; not helpful
        # Better: use equation mu*v = I + C directly:
        C_value_raw = (mu * v0 - I.subs(x, parsed_x0))
        C_value = sp.simplify(C_value_raw) if simplify else C_value_raw

        v_particular = (I + C_value) / mu
        if simplify:
            v_particular = sp.simplify(v_particular)

        y_particular_raw = sp.Pow(v_particular, power_inv)
        y_particular = sp.simplify(y_particular_raw) if simplify else y_particular_raw

    # Notes (optional)
    note_parts = []
    if bool(intP_lin.has(sp.Integral)):
        note_parts.append("SymPy kept an integral unevaluated in the integrating factor.")
    if bool(I.has(sp.Integral)):
        note_parts.append("SymPy kept an integral unevaluated in the solution integral.")
    anti_note = " ".join(note_parts)

    # -----------------------------
    # 3) Build step-by-step HTML
    # -----------------------------
    ode_ltx = f"\\frac{{d {L(yx)}}}{{d {L(x)}}} + \\left({L(P_)}\\right){L(yx)} = \\left({L(Q_)}\\right){L(yx)}^{{{L(n_)}}}"

    # Multiply by y^{-n}
    step_mul_ltx = f"{L(yx)}^{{-{L(n_)}}}\\frac{{d {L(yx)}}}{{d {L(x)}}} + \\left({L(P_)}\\right){L(yx)}^{{{L(1-n_)}}} = {L(Q_)}"

    v_def_ltx = f"v = {L(yx)}^{{{L(1-n_)}}}"
    vprime_rel_ltx = f"\\frac{{dv}}{{d{L(x)}}} = \\left({L(1-n_)}\\right){L(yx)}^{{-{L(n_)}}}\\frac{{d {L(yx)}}}{{d {L(x)}}}"

    # Replace y^{-n} y' = v'/(1-n)
    lin_in_v_ltx = f"\\frac{{1}}{{{L(1-n_)}}}\\frac{{dv}}{{d{L(x)}}} + \\left({L(P_)}\\right)v = {L(Q_)}"
    lin_scaled_ltx = f"\\frac{{dv}}{{d{L(x)}}} + \\left({L(P_lin)}\\right)v = {L(Q_lin)} \\qquad \\cdots \\text{{(Linear in }}v\\text{{)}}"

    mu_def_ltx = f"\\mu({L(x)}) = e^{{\\int {L(P_lin)}\\, d{L(x)}}}"
    mu_val_ltx = f"\\mu({L(x)}) = e^{{{L(intP_lin)}}} = {L(mu)}"

    deriv_ltx = f"\\frac{{d}}{{d{L(x)}}}\\left({L(mu)}\\,v\\right) = {L(mu)}\\left({L(Q_lin)}\\right)"
    integ_ltx = f"{L(mu)}\\,v = \\int {L(mu)}\\left({L(Q_lin)}\\right)\\, d{L(x)} + {L(C)}"
    I_ltx = f"\\int {L(mu)}\\left({L(Q_lin)}\\right)\\, d{L(x)} = {L(I)}"
    v_sol_ltx = f"v = \\frac{{{L(I)} + {L(C)}}}{{{L(mu)}}}"
    y_from_v_ltx = f"{L(yx)} = v^{{\\frac{{1}}{{{L(1-n_)}}}}} = \\left(\\frac{{{L(I)} + {L(C)}}}{{{L(mu)}}}\\right)^{{\\frac{{1}}{{{L(1-n_)}}}}}"

    soln_stmts = []

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

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

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

        & \\displaystyle \\Rightarrow \\text{{Multiply both sides by }} {L(yx)}^{{-{L(n_)}}}: \\\\
        & \\displaystyle \\Rightarrow {step_mul_ltx} \\\\[6pt]

        & \\displaystyle \\Rightarrow \\text{{Let }} {v_def_ltx} \\\\
        & \\displaystyle \\Rightarrow {vprime_rel_ltx} \\\\[6pt]

        & \\displaystyle \\Rightarrow {lin_in_v_ltx} \\\\
        & \\displaystyle \\Rightarrow {lin_scaled_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow \\text{{Integrating factor: }} {mu_def_ltx} \\\\
        & \\displaystyle \\Rightarrow {mu_val_ltx} \\\\[6pt]

        & \\displaystyle \\Rightarrow {deriv_ltx} \\\\
        & \\displaystyle \\Rightarrow {integ_ltx} \\\\
        & \\displaystyle \\Rightarrow {I_ltx} \\\\[6pt]

        & \\displaystyle \\Rightarrow {v_sol_ltx} \\\\
        & \\displaystyle \\Rightarrow {y_from_v_ltx}  \\qquad \\cdots \\text{{(i)}} \\\\

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

    ic_ltx = ""
    final_ic_ltx = ""
    if use_ic and (y_particular is not None) and (C_value is not None):
        ic_ltx = f"{L(yx)}\\Big|_{{{L(x)}={L(parsed_x0)}}} = {L(parsed_y0)}"
        v0_ltx = f"v\\Big|_{{{L(x)}={L(parsed_x0)}}} = {L(parsed_y0)}^{{{L(1-n_)}}} = {L(v0)}"

        # C = mu*v - I at x0
        C_expr_ltx = f"{L(C)} = {L(mu)}\\,v - {L(I)}"
        C_at_x0_ltx = (
            f"{L(C)} = {L(mu.subs(x, parsed_x0))}\\,{L(v0)} - {L(I.subs(x, parsed_x0))}"
            f" \\qquad \\text{{at }} {L(x)}={L(parsed_x0)}"
        )
        C_val_ltx = f"{L(C)} = {L(C_value)} \\qquad \\text{{Given }} {ic_ltx}"

        final_ic_ltx = f"{L(yx)} = {L(y_particular)} "

        ic_block = f"""
            <br/>
            \\(\\displaystyle
            \\begin{{aligned}}
            & \\displaystyle \\Rightarrow {v0_ltx} \\\\
            & \\displaystyle \\Rightarrow {C_expr_ltx} \\\\
            & \\displaystyle \\Rightarrow {C_at_x0_ltx} \\\\
            & \\displaystyle \\Rightarrow {C_val_ltx} \\\\
            & \\displaystyle \\Rightarrow {final_ic_ltx}  \\qquad \\cdots \\text{{ from (i)}} \\\\
            \\end{{aligned}}
            \\)
        """
        soln_stmts.append(ic_block.replace("\n", ""))

    # -----------------------------
    # 4) Final problem/answer strings
    # -----------------------------
    if use_ic:
        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {ic_ltx} \\)"
        answer_s = f"\\( {final_ic_ltx} \\)" if final_ic_ltx else f"\\( {L(yx)} = {L(y_general)} \\)"
    else:
        problem_s = f"Solve \\( {ode_ltx} \\)"
        answer_s = f"\\( {L(yx)} = {L(y_general)} \\)"

    soln_s = "".join(soln_stmts)

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


