# ============================================================
# File: _solve_ode_first_order_exact_integrating_factor.py
# Topic: First-Order ODE (Non-Exact -> Exact via Integrating Factor)
#
# Form:   M(x,y) dx + N(x,y) dy = 0
#
# Exactness condition:
#   ∂M/∂y == ∂N/∂x
#
# If NOT exact, try integrating factor:
#   A) If  ( (∂M/∂y - ∂N/∂x) / N ) = f(x) only,
#         μ(x) = exp( ∫ f(x) dx )
#
#   B) If  ( (∂N/∂x - ∂M/∂y) / M ) = g(y) only,
#         μ(y) = exp( ∫ g(y) dy )
#
# Then μM dx + μN dy = 0 becomes exact.
# Potential function Ψ(x,y):
#   Ψ_x = μM,  Ψ_y = μN
# Implicit solution:  Ψ(x,y) = C
# Optional IC: y(x0)=y0 => C = Ψ(x0,y0)
#
# Notes:
#   - If SymPy can’t integrate, it keeps Integral(...)
#   - We primarily return an implicit solution Ψ(x,y)=C
# ============================================================

from __future__ import annotations

from typing import Optional, Union, Dict, Any

import sympy as sp


def solve_ode_first_order_exact_integrating_factor(
    M_of_xy: Union[str, sp.Expr],
    N_of_xy: 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:
        M(x,y) dx + N(x,y) dy = 0

    If not exact, attempt integrating factor μ(x) or μ(y).

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

    Output style:
      - HTML with \\(\\displaystyle ...\\) and aligned steps
      - "answer_s" is the final implicit solution (no "from (i)" tags)
    """

    # -----------------------------
    # 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)
    C = sp.Symbol(C_symbol)

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

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

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

    M = sp.simplify(M_raw) if simplify else M_raw
    N = sp.simplify(N_raw) if simplify else N_raw

    use_ic = (x0 is not None) and (y0 is not None)
    parsed_x0 = parsed_y0 = None
    if use_ic:
        parsed_x0 = sp.sympify(x0, locals=base_locals)
        parsed_y0 = sp.sympify(y0, locals=base_locals)

    # -----------------------------
    # 2) Check exactness
    # -----------------------------
    My = sp.diff(M, yy)
    Nx = sp.diff(N, x)
    exact_test = sp.simplify(My - Nx) if simplify else (My - Nx)
    is_exact = bool(exact_test == 0)

    # -----------------------------
    # 3) Find integrating factor μ(x) or μ(y) if needed
    # -----------------------------
    mu = sp.Integer(1)
    mu_kind = None   # "x" or "y" or None
    fx = None
    gy = None

    if not is_exact:
        # Try μ(x): f(x) = (My - Nx)/N depends only on x
        try:
            fx_raw = sp.simplify((My - Nx) / N) if simplify else (My - Nx) / N
            if fx_raw.free_symbols.issubset({x}):
                fx = sp.simplify(fx_raw) if simplify else fx_raw
                mu = sp.exp(sp.integrate(fx, x))
                mu_kind = "x"
        except Exception:
            pass

        # Try μ(y): g(y) = (Nx - My)/M depends only on y
        if mu_kind is None:
            try:
                gy_raw = sp.simplify((Nx - My) / M) if simplify else (Nx - My) / M
                if gy_raw.free_symbols.issubset({yy}):
                    gy = sp.simplify(gy_raw) if simplify else gy_raw
                    mu = sp.exp(sp.integrate(gy, yy))
                    mu_kind = "y"
            except Exception:
                pass

    # If still not found and not exact, return a graceful result
    if (not is_exact) and (mu_kind is None):
        ode_ltx = f"\\left({L(M.subs(yy, yx))}\\right)\\,d{L(x)} + \\left({L(N.subs(yy, yx))}\\right)\\,d{L(yx)} = 0"
        problem_s = f"Solve \\( {ode_ltx} \\) using integrating factor (if possible)."
        soln_s = (
            f"<div class='wrap'><div class='card'>"
            f"<div class='badge'>Note: Non-Exact ODE → Integrating Factor (μ(x) or μ(y))<br/></div>"
            f"\\(\\displaystyle \\begin{{aligned}}"
            f"& \\displaystyle M={L(M)},\\; N={L(N)} \\\\"
            f"& \\displaystyle \\frac{{\\partial M}}{{\\partial y}}={L(My)},\\; \\frac{{\\partial N}}{{\\partial x}}={L(Nx)} \\\\"
            f"& \\displaystyle \\Rightarrow \\text{{Not exact since }} \\frac{{\\partial M}}{{\\partial y}}-\\frac{{\\partial N}}{{\\partial x}}={L(exact_test)} \\neq 0 \\\\"
            f"& \\displaystyle \\Rightarrow \\text{{No simple }}\\mu(x)\\text{{ or }}\\mu(y)\\text{{ detected by the standard tests.}} \\\\"
            f"\\end{{aligned}}\\)"
            f"</div></div>"
        )
        answer_s = "\\(\\text{No μ(x) or μ(y) found by the standard tests.}\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # -----------------------------
    # 4) Apply μ and form exact equation
    # -----------------------------
    M1 = sp.simplify(mu * M) if simplify else mu * M
    N1 = sp.simplify(mu * N) if simplify else mu * N

    # Verify exactness after μ (best-effort)
    My1 = sp.diff(M1, yy)
    Nx1 = sp.diff(N1, x)
    exact_test2 = sp.simplify(My1 - Nx1) if simplify else (My1 - Nx1)

    # -----------------------------
    # 5) Potential function Ψ(x,y)
    # -----------------------------
    # Ψ_x = M1 -> Ψ = ∫ M1 dx + g(y)
    Psi_part = sp.integrate(M1, x)
    gprime = sp.simplify(N1 - sp.diff(Psi_part, yy)) if simplify else (N1 - sp.diff(Psi_part, yy))
    g_of_y = sp.integrate(gprime, yy)
    Psi = sp.simplify(Psi_part + g_of_y) if simplify else (Psi_part + g_of_y)

    # IC constant
    C_value = None
    if use_ic:
        try:
            C_value_raw = Psi.subs({x: parsed_x0, yy: parsed_y0})
            C_value = sp.simplify(C_value_raw) if simplify else C_value_raw
        except Exception:
            C_value = None

    # Substitute yy -> y(x) for display
    Psi_yx = Psi.subs(yy, yx)

    # -----------------------------
    # 6) Build HTML (steps) - match your style
    # -----------------------------
    ode_ltx = f"\\left({L(M.subs(yy, yx))}\\right)\\,d{L(x)} + \\left({L(N.subs(yy, yx))}\\right)\\,d{L(yx)} = 0"
    exact1_ltx = f"\\frac{{\\partial M}}{{\\partial y}}={L(My)},\\; \\frac{{\\partial N}}{{\\partial x}}={L(Nx)}"
    exact2_ltx = f"\\frac{{\\partial M}}{{\\partial y}}-\\frac{{\\partial N}}{{\\partial x}}={L(exact_test)}"

    note_parts = []
    if mu_kind == "x":
        note_parts.append("Integrating factor detected as \\(\\mu(x)\\).")
    elif mu_kind == "y":
        note_parts.append("Integrating factor detected as \\(\\mu(y)\\).")
    if bool(mu.has(sp.Integral)):
        note_parts.append("SymPy kept an integral inside \\(\\mu\\) unevaluated.")
    if bool(Psi.has(sp.Integral)):
        note_parts.append("SymPy kept an integral inside \\(\\Psi\\) unevaluated.")
    if exact_test2 != 0:
        # still show it; sometimes simplification doesn't reduce fully
        note_parts.append("Exactness after multiplying by \\(\\mu\\) could not be simplified to 0, but we proceed symbolically.")

    anti_note = " ".join(note_parts)

    # μ display
    if mu_kind == "x":
        mu_rule_ltx = f"\\mu(x)=\\exp\\left(\\int {L(fx)}\\,d{L(x)}\\right) = {L(mu)}"
    elif mu_kind == "y":
        mu_rule_ltx = f"\\mu(y)=\\exp\\left(\\int {L(gy)}\\,d{L(yy)}\\right) = {L(mu)}"
    else:
        mu_rule_ltx = f"\\mu=1"

    exact_after_ltx = f"\\tilde M=\\mu M={L(M1)},\\; \\tilde N=\\mu N={L(N1)}"
    exact_after_test_ltx = f"\\frac{{\\partial \\tilde M}}{{\\partial y}}-\\frac{{\\partial \\tilde N}}{{\\partial x}}={L(exact_test2)}"

    psi_build_ltx = (
        f"\\Psi_x=\\tilde M \\Rightarrow \\Psi=\\int \\tilde M\\,d{L(x)}+g(y)"
    )
    gprime_ltx = f"g'(y)=\\tilde N-\\frac{{\\partial}}{{\\partial y}}\\left(\\int \\tilde M\\,d{L(x)}\\right) = {L(gprime)}"
    psi_final_ltx = f"\\Psi(x,y)={L(Psi)}"

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: Non-Exact ODE → Integrating Factor (μ(x) or μ(y))<br/></div>

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

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

        & \\displaystyle \\Rightarrow \\text{{Check exactness:}} \\\\
        & \\displaystyle \\Rightarrow {exact1_ltx} \\\\
        & \\displaystyle \\Rightarrow {exact2_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Find an integrating factor:}} \\\\
        & \\displaystyle \\Rightarrow {mu_rule_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Multiply the equation by }}\\mu: \\\\
        & \\displaystyle \\Rightarrow {exact_after_ltx} \\\\
        & \\displaystyle \\Rightarrow {exact_after_test_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Potential function }}\\Psi(x,y): \\\\
        & \\displaystyle \\Rightarrow {psi_build_ltx} \\\\
        & \\displaystyle \\Rightarrow {gprime_ltx} \\\\
        & \\displaystyle \\Rightarrow {psi_final_ltx} \\\\[8pt]

        & \\displaystyle \\Rightarrow \\text{{Implicit solution: }} \\Psi(x,y)=C \\\\

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

    # IC block
    if use_ic and (C_value is not None):
        ic_ltx = f"{L(yx)}\\Big|_{{{L(x)}={L(parsed_x0)}}} = {L(parsed_y0)}"
        C_val_ltx = f"{L(C)} = \\Psi({L(parsed_x0)},{L(parsed_y0)}) = {L(C_value)} \\qquad \\text{{Given }} {ic_ltx}"
        final_soln_ltx = f"{L(Psi_yx)} = {L(C_value)}"
        final_ans_ltx = f"{L(Psi_yx)} = {L(C_value)}"

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

        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {ic_ltx} \\)"
        answer_s = f"\\( {final_ans_ltx} \\)"
    else:
        problem_s = f"Solve \\( {ode_ltx} \\)"
        answer_s = f"\\( {L(Psi_yx)} = {L(C)} \\)"

    soln_s = "".join(soln_stmts)

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





