# ============================================================
# File: _solve_ode_first_order_separable.py
# Topic: First-Order Separable ODE
#
# Form:
#     dy/dx = g(x) * h(y)
#
# Method:
#     dy/h(y) = g(x) dx
#     ∫ dy/h(y) = ∫ g(x) dx
#     F(y) = G(x) + C
#
# Optional IC:
#     y(x0) = y0  =>  C = F(y0) - G(x0)
#
# Notes:
#   - If SymPy can't integrate, Integral(...) is kept.
#   - We keep an implicit solution; explicit y(x) is attempted best-effort.
# ============================================================

from __future__ import annotations

from typing import Any, Dict, Optional, Union

import sympy as sp


def solve_ode_first_order_separable(
    g_of_x: Union[str, sp.Expr],
    h_of_y: 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 the separable first-order ODE:
        dy/dx = g(x) * h(y)

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

    Output style:
      - HTML with \\(\\displaystyle ...\\) and aligned steps
      - "answer_s" is the final implicit (or explicit if found) solution
    """

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

    yy = sp.Symbol("y", real=True)
    C = sp.Symbol(C_symbol)

    base_locals = {
        "x": x,
        "y": yy,  # for parsing h(y)
        "E": sp.E,
        "pi": sp.pi,
        "sin": sp.sin,
        "cos": sp.cos,
        "tan": sp.tan,
        "exp": sp.exp,
        "log": sp.log,
        "ln": sp.log,
        "sqrt": sp.sqrt,
        "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:
        g_raw = sp.sympify(g_of_x, locals=base_locals) if isinstance(g_of_x, str) else sp.sympify(g_of_x)
    except Exception as e:
        raise ValueError(f"Could not parse g(x)={g_of_x!r}. Error: {e}") from e

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

    g = sp.simplify(g_raw) if simplify else g_raw
    h = sp.simplify(h_raw) if simplify else h_raw

    # Parse ICs (best-effort)
    parsed_x0 = None
    parsed_y0 = None
    if x0 is not None and y0 is not None:
        try:
            parsed_x0 = sp.sympify(x0, locals=base_locals)
            parsed_y0 = sp.sympify(y0, locals=base_locals)
        except Exception:
            parsed_x0 = x0
            parsed_y0 = y0

    # -----------------------------
    # 2) Integrate both sides
    # -----------------------------
    # Left: ∫ dy / h(y)
    left_integrand = sp.simplify(1 / h) if simplify else 1 / h
    F_raw = sp.integrate(left_integrand, yy)
    F = sp.simplify(F_raw) if simplify else F_raw

    # Right: ∫ g(x) dx
    G_raw = sp.integrate(g, x)
    G = sp.simplify(G_raw) if simplify else G_raw

    # If integration fails, SymPy often returns an Integral automatically; that is OK.

    # -----------------------------
    # 3) Optional IC -> determine C
    # -----------------------------
    C_value = None
    if parsed_x0 is not None and parsed_y0 is not None:
        try:
            C_value_raw = F.subs(yy, parsed_y0) - G.subs(x, parsed_x0)
            C_value = sp.simplify(C_value_raw) if simplify else C_value_raw
        except Exception:
            C_value = None

    # -----------------------------
    # 4) Build final implicit / explicit
    # -----------------------------
    # Implicit: F(y(x)) = G(x) + C  (or + C_value)
    F_yx = F.subs(yy, yx)

    if C_value is not None:
        implicit_eq = sp.Eq(F_yx, G + C_value)
        implicit_ans_ltx = f"{L(F_yx)} = {L(G)} + {L(C_value)}"
    else:
        implicit_eq = sp.Eq(F_yx, G + C)
        implicit_ans_ltx = f"{L(F_yx)} = {L(G)} + {L(C)}"

    # Try explicit y(x) best-effort
    explicit_expr = None
    try:
        # Solve F(yy) = G(x) + const for yy, then substitute yy -> y(x)
        target_rhs = (G + (C_value if C_value is not None else C))
        sol_yy = sp.solve(sp.Eq(F, target_rhs), yy)
        if sol_yy:
            explicit_expr = sol_yy[0].subs(yy, yx)
    except Exception:
        explicit_expr = None

    # -----------------------------
    # 5) HTML formatting (match your style)
    # -----------------------------
    ode_ltx = f"\\frac{{d {L(yx)}}}{{d {L(x)}}} = \\left({L(g)}\\right)\\left({L(h.subs(yy, yx))}\\right)"
    sep_ltx = f"\\frac{{d{L(yx)}}}{{{L(h.subs(yy, yx))}}} = \\left({L(g)}\\right)d{L(x)}"
    int_ltx = f"\\int \\frac{{d{L(yx)}}}{{{L(h.subs(yy, yx))}}} = \\int \\left({L(g)}\\right)d{L(x)}"

    left_int_ltx = f"\\int \\frac{{dy}}{{{L(h)}}}"
    right_int_ltx = f"\\int {L(g)}\\,dx"

    Fy_ltx = L(F) if not F.has(sp.Integral) else L(sp.Integral(left_integrand, yy))
    Gx_ltx = L(G) if not G.has(sp.Integral) else L(sp.Integral(g, x))

    # Step card
    soln_stmts = []

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

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

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

        & \\displaystyle \\Rightarrow {sep_ltx} \\\\[6pt]

        & \\displaystyle \\Rightarrow {int_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow {Fy_ltx} = {Gx_ltx} + {L(C)} \\qquad \\cdots (i) \\\\[6pt]

        \\end{{aligned}}
        \\)

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

    # IC block (if available + computed)
    if parsed_x0 is not None and parsed_y0 is not None and C_value is not None:
        ic_ltx = f"{L(y)}\\left({L(parsed_x0)}\\right)={L(parsed_y0)}"
        C_expr_ltx = f"{L(C)} = {Fy_ltx}\\Big|_{{y={L(parsed_y0)}}} - {Gx_ltx}\\Big|_{{x={L(parsed_x0)}}}"
        C_val_ltx = f"{L(C)} = {L(C_value)} \\qquad \\text{{Given }} {ic_ltx}"
        final_soln_ltx = f"{implicit_ans_ltx} \\qquad \\cdots \\text{{from (i)}}"

        ic_block = f"""
            <br/>
            \\(\\displaystyle
            \\begin{{aligned}}
            & \\displaystyle \\Rightarrow {C_expr_ltx} \\\\
            & \\displaystyle \\Rightarrow {C_val_ltx} \\\\
            & \\displaystyle \\Rightarrow {final_soln_ltx} \\\\
            \\end{{aligned}}
            \\)
        """
        soln_stmts.append(ic_block.replace("\n", ""))

    # Note if integrals unevaluated
    note_parts = []
    if F.has(sp.Integral) or G.has(sp.Integral):
        note_parts.append("SymPy kept one or both integrals unevaluated.")

    if note_parts:
        note_html = f"""
        <div class="note">
          <b>Note.</b> {' '.join(note_parts)}
        </div>
        """
        soln_stmts.append(note_html.replace("\n", ""))

    # Problem + answer
    if parsed_x0 is not None and parsed_y0 is not None:
        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {L(y)}({L(parsed_x0)})={L(parsed_y0)} \\)"
    else:
        problem_s = f"Solve \\( {ode_ltx} \\)"

    if explicit_expr is not None:
        answer_s = f"\\( {L(yx)} = {L(explicit_expr)} \\)"
    else:
        answer_s = f"\\( {implicit_ans_ltx} \\)"

    soln_s = "".join(soln_stmts)

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