# ============================================================
# File: _solve_ode_second_order_variation_of_parameters.py
# Topic: 2nd-Order Linear Nonhomogeneous ODE (Variation of Parameters)
#
# Form:
#   y'' + P(x) y' + Q(x) y = R(x)
#
# Given a fundamental set {y1(x), y2(x)} of the homogeneous equation:
#   y'' + P(x) y' + Q(x) y = 0
#
# Wronskian:
#   W = y1*y2' - y1'*y2
#
# Variation of parameters:
#   u1' = - y2*R / W
#   u2' =   y1*R / W
#   y_p = u1*y1 + u2*y2
#
# General solution:
#   y = C1*y1 + C2*y2 + y_p
#
# Notes:
#   - If SymPy can’t integrate, it keeps Integral(...)
# ============================================================

from __future__ import annotations

from typing import Any, Dict, Optional, Union

import sympy as sp


def solve_ode_second_order_variation_of_parameters(
    P_of_x: Union[str, sp.Expr],
    Q_of_x: Union[str, sp.Expr],
    R_of_x: Union[str, sp.Expr],
    y1_of_x: Union[str, sp.Expr],
    y2_of_x: Union[str, sp.Expr],
    *,
    x_symbol: str = "x",
    y_symbol: str = "y",
    C1_symbol: str = "C_1",
    C2_symbol: str = "C_2",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Solve y'' + P(x) y' + Q(x) y = R(x) using variation of parameters,
    given y1(x), y2(x) for the associated homogeneous equation.

    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)

    C1 = sp.Symbol(C1_symbol)
    C2 = sp.Symbol(C2_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)

    P_raw = _parse(P_of_x)
    Q_raw = _parse(Q_of_x)
    R_raw = _parse(R_of_x)
    y1_raw = _parse(y1_of_x)
    y2_raw = _parse(y2_of_x)

    P = sp.simplify(P_raw) if simplify else P_raw
    Q = sp.simplify(Q_raw) if simplify else Q_raw
    R = sp.simplify(R_raw) if simplify else R_raw
    y1 = sp.simplify(y1_raw) if simplify else y1_raw
    y2 = sp.simplify(y2_raw) if simplify else y2_raw

    # -----------------------------
    # 2) Best-effort validate y1,y2 solve homogeneous
    # -----------------------------
    res1 = None
    res2 = None
    try:
        res1_raw = sp.diff(y1, x, 2) + P * sp.diff(y1, x) + Q * y1
        res1 = sp.simplify(res1_raw) if simplify else res1_raw
    except Exception:
        res1 = None

    try:
        res2_raw = sp.diff(y2, x, 2) + P * sp.diff(y2, x) + Q * y2
        res2 = sp.simplify(res2_raw) if simplify else res2_raw
    except Exception:
        res2 = None

    # -----------------------------
    # 3) Compute Wronskian and u1', u2'
    # -----------------------------
    W = sp.simplify(y1 * sp.diff(y2, x) - sp.diff(y1, x) * y2) if simplify else (y1 * sp.diff(y2, x) - sp.diff(y1, x) * y2)

    # If W=0, we cannot proceed
    if simplify:
        if sp.simplify(W) == 0:
            raise ValueError("Wronskian W=0. y1,y2 are not linearly independent.")
    else:
        if W == 0:
            raise ValueError("Wronskian W=0. y1,y2 are not linearly independent.")

    u1_prime = sp.simplify(-y2 * R / W) if simplify else (-y2 * R / W)
    u2_prime = sp.simplify(y1 * R / W) if simplify else (y1 * R / W)

    u1 = sp.integrate(u1_prime, x)
    u2 = sp.integrate(u2_prime, x)

    y_p = sp.simplify(u1 * y1 + u2 * y2) if simplify else (u1 * y1 + u2 * y2)
    y_general = sp.simplify(C1 * y1 + C2 * y2 + y_p) if simplify else (C1 * y1 + C2 * y2 + y_p)

    # -----------------------------
    # 4) HTML formatting (match your style)
    # -----------------------------
    ode_ltx = (
        f"\\frac{{d^2{L(yx)}}}{{d{L(x)}^2}}"
        f" + \\left({L(P)}\\right)\\frac{{d{L(yx)}}}{{d{L(x)}}}"
        f" + \\left({L(Q)}\\right){L(yx)}"
        f" = {L(R)}"
    )

    y1_ltx = f"y_1(x)={L(y1)}"
    y2_ltx = f"y_2(x)={L(y2)}"
    W_ltx = L(W)

    u1p_ltx = L(u1_prime)
    u2p_ltx = L(u2_prime)

    u1_ltx = L(u1) if not u1.has(sp.Integral) else L(sp.Integral(u1_prime, x))
    u2_ltx = L(u2) if not u2.has(sp.Integral) else L(sp.Integral(u2_prime, x))

    yp_ltx = L(y_p) if not y_p.has(sp.Integral) else f"{L(y1)}\\left({u1_ltx}\\right)+{L(y2)}\\left({u2_ltx}\\right)"

    note_parts = []
    if res1 is not None and res1 != 0:
        note_parts.append("Given \\(y_1\\) did not simplify to an exact homogeneous solution; proceeding symbolically.")
    if res2 is not None and res2 != 0:
        note_parts.append("Given \\(y_2\\) did not simplify to an exact homogeneous solution; proceeding symbolically.")
    if (u1.has(sp.Integral) or u2.has(sp.Integral) or y_p.has(sp.Integral)):
        note_parts.append("SymPy kept one or more integrals unevaluated.")

    anti_note = " ".join(note_parts)

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: 2nd-Order Linear Nonhomogeneous ODE → Variation of Parameters<br/></div>

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

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

        & \\displaystyle \\text{{Given a fundamental set for the homogeneous equation:}} \\\\
        & \\displaystyle {y1_ltx},\\quad {y2_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow W = y_1y_2' - y_1'y_2 = {W_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow u_1'(x) = -\\frac{{y_2 R}}{{W}} = {u1p_ltx} \\\\
        & \\displaystyle \\Rightarrow u_2'(x) = \\frac{{y_1 R}}{{W}} = {u2p_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow u_1(x)=\\int u_1'(x)\\,dx = {u1_ltx} \\\\
        & \\displaystyle \\Rightarrow u_2(x)=\\int u_2'(x)\\,dx = {u2_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow y_p = u_1y_1 + u_2y_2 = {yp_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{General solution: }} y = C_1y_1 + C_2y_2 + y_p \\\\
        & \\displaystyle \\Rightarrow {L(yx)} = {L(y_general)} \\\\

        \\end{{aligned}}
        \\)

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

    problem_s = (
        f"Solve \\( {ode_ltx} \\) by variation of parameters, "
        f"given homogeneous solutions \\( {y1_ltx} \\), \\( {y2_ltx} \\)."
    )
    answer_s = f"\\( y_p={yp_ltx},\\quad {L(yx)} = {L(y_general)} \\)"

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