# ============================================================
# File: _solve_ode_second_order_linear_nonhomogeneous_const_coeff.py
# Topic: 2nd-order linear nonhomogeneous ODE (constant coefficients)
# Form:  a y'' + b y' + c y = g(x),  a != 0
# Method:
#   1) Complementary solution y_c from characteristic equation
#   2) Particular solution y_p:
#        - try Undetermined Coefficients (UC) via SymPy hint
#        - fallback to Variation of Parameters (VoP) via SymPy hint
#   3) y = y_c + y_p
#   4) Optional IC: y(x0)=y0, y'(x0)=dy0
# Output: dict {problem_s, soln_s, answer_s}
# ============================================================

from __future__ import annotations

from typing import Optional, Union, Dict, Any, Tuple

import sympy as sp


def solve_ode_second_order_linear_nonhomogeneous_const_coeff(
    a: Union[str, sp.Expr],
    b: Union[str, sp.Expr],
    c: Union[str, sp.Expr],
    g: Union[str, sp.Expr],
    *,
    x_symbol: str = "x",
    y_symbol: str = "y",
    r_symbol: str = "r",
    C1_symbol: str = "C_1",
    C2_symbol: str = "C_2",
    x0: Optional[Union[int, float, str, sp.Expr]] = None,
    y0: Optional[Union[int, float, str, sp.Expr]] = None,
    dy0: Optional[Union[int, float, str, sp.Expr]] = None,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Full-steps HTML solution for:
        a y'' + b y' + c y = g(x),   with constant coefficients a,b,c.

    Optional initial conditions:
        y(x0) = y0
        y'(x0) = dy0

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

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

    r = sp.Symbol(r_symbol)
    C1 = sp.Symbol(C1_symbol)
    C2 = sp.Symbol(C2_symbol)

    base_locals = {
        x_symbol: x,
        r_symbol: r,
        "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:
        a_ = sp.sympify(a, locals=base_locals) if isinstance(a, str) else sp.sympify(a)
        b_ = sp.sympify(b, locals=base_locals) if isinstance(b, str) else sp.sympify(b)
        c_ = sp.sympify(c, locals=base_locals) if isinstance(c, str) else sp.sympify(c)
        g_ = sp.sympify(g, locals=base_locals) if isinstance(g, str) else sp.sympify(g)
    except Exception as e:
        raise ValueError(f"Could not parse inputs. Error: {e}") from e

    if simplify:
        a_, b_, c_, g_ = sp.simplify(a_), sp.simplify(b_), sp.simplify(c_), sp.simplify(g_)

    if a_ == 0:
        raise ValueError("Coefficient a must be non-zero for a second-order ODE.")

    # -----------------------------
    # 2) Complementary solution y_c
    # -----------------------------
    D = sp.simplify(b_**2 - 4*a_*c_) if simplify else (b_**2 - 4*a_*c_)
    r1 = sp.simplify((-b_ + sp.sqrt(D)) / (2*a_)) if simplify else ((-b_ + sp.sqrt(D)) / (2*a_))
    r2 = sp.simplify((-b_ - sp.sqrt(D)) / (2*a_)) if simplify else ((-b_ - sp.sqrt(D)) / (2*a_))

    # Decide root type if possible
    case = "unknown"
    if D.is_zero is True:
        case = "repeated"
    elif D.is_positive is True:
        case = "distinct_real"
    elif D.is_negative is True:
        case = "complex"

    if case == "repeated":
        r0 = r1
        yc = (C1 + C2*x) * sp.exp(r0*x)
        yc_note = "D=0 (repeated root)"
    elif case == "complex":
        alpha = sp.simplify(-b_/(2*a_)) if simplify else (-b_/(2*a_))
        beta = sp.simplify(sp.sqrt(-D)/(2*a_)) if simplify else (sp.sqrt(-D)/(2*a_))
        yc = sp.exp(alpha*x) * (C1*sp.cos(beta*x) + C2*sp.sin(beta*x))
        yc_note = "D<0 (complex roots)"
    else:
        yc = C1*sp.exp(r1*x) + C2*sp.exp(r2*x)
        yc_note = "use r1, r2"

    if simplify:
        yc = sp.simplify(yc)

    # -----------------------------
    # 3) Particular solution y_p (UC preferred, VoP fallback)
    # -----------------------------
    ode = sp.Eq(a_*sp.diff(yx, x, 2) + b_*sp.diff(yx, x) + c_*yx, g_)

    yp = None
    method_used = None

    # Try undetermined coefficients hint
    try:
        sol_uc = sp.dsolve(ode, yx, hint="nth_linear_constant_coeff_undetermined_coefficients")
        # sol_uc.rhs is full y(x), includes constants -> strip to yp by subtracting yc form?
        # Easier: compute yp by plugging a particular form from the solution and subtracting a homogeneous part:
        # We'll extract yp by setting C1=C2=0 in the dsolve result (works for these hints).
        y_full_uc = sol_uc.rhs
        yp_uc = y_full_uc.subs({sp.Symbol("C1"): 0, sp.Symbol("C2"): 0, C1: 0, C2: 0})
        yp = sp.simplify(yp_uc) if simplify else yp_uc
        method_used = "Undetermined Coefficients"
    except Exception:
        pass

    if yp is None:
        # Fallback: variation of parameters
        try:
            sol_vop = sp.dsolve(ode, yx, hint="nth_linear_constant_coeff_variation_of_parameters")
            y_full_vop = sol_vop.rhs
            yp_vop = y_full_vop.subs({sp.Symbol("C1"): 0, sp.Symbol("C2"): 0, C1: 0, C2: 0})
            yp = sp.simplify(yp_vop) if simplify else yp_vop
            method_used = "Variation of Parameters"
        except Exception as e:
            raise ValueError(f"Could not compute a particular solution (UC and VoP failed). Error: {e}") from e

    y_general = yc + yp
    if simplify:
        y_general = sp.simplify(y_general)

    # -----------------------------
    # 4) Optional initial conditions
    # -----------------------------
    use_ic = (x0 is not None) and (y0 is not None) and (dy0 is not None)

    y_particular = None
    C1_value = None
    C2_value = None

    if use_ic:
        x0_ = sp.sympify(x0, locals=base_locals)
        y0_ = sp.sympify(y0, locals=base_locals)
        dy0_ = sp.sympify(dy0, locals=base_locals)

        yprime_general = sp.diff(y_general, x)

        eq1 = sp.Eq(y_general.subs(x, x0_), y0_)
        eq2 = sp.Eq(yprime_general.subs(x, x0_), dy0_)

        solC = sp.solve((eq1, eq2), (C1, C2), dict=True)
        if solC:
            C1_value = solC[0].get(C1, None)
            C2_value = solC[0].get(C2, None)
            y_particular = y_general.subs({C1: C1_value, C2: C2_value})
            if simplify:
                C1_value = sp.simplify(C1_value) if C1_value is not None else None
                C2_value = sp.simplify(C2_value) if C2_value is not None else None
                y_particular = sp.simplify(y_particular)

    # -----------------------------
    # 5) Build step-by-step HTML
    # -----------------------------
    y1 = sp.diff(yx, x)
    y2 = sp.diff(yx, x, 2)

    ode_ltx = f"{L(a_)}\\,{L(y2)} + {L(b_)}\\,{L(y1)} + {L(c_)}\\,{L(yx)} = {L(g_)}"

    char_ltx = f"{L(a_)}{L(r)}^2 + {L(b_)}{L(r)} + {L(c_)} = 0"
    D_ltx = f"D = {L(b_)}^2 - 4({L(a_)})({L(c_)}) = {L(D)}"
    r1_ltx = f"{L(r)}_1 = \\frac{{-{L(b_)} + \\sqrt{{{L(D)}}}}}{{2({L(a_)})}} = {L(r1)}"
    r2_ltx = f"{L(r)}_2 = \\frac{{-{L(b_)} - \\sqrt{{{L(D)}}}}}{{2({L(a_)})}} = {L(r2)}"
    yc_ltx = f"{L(yx)}_c = {L(yc)} \\qquad \\cdots \\text{{(i)}}"

    yp_ltx = f"{L(yx)}_p = {L(yp)} \\qquad \\cdots \\text{{(ii)}}"
    ygen_ltx = f"{L(yx)} = {L(y_general)} \\qquad \\cdots \\text{{(iii)}}"

    method_note = f"Method for particular solution: {method_used}"

    soln_stmts = []

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: Second-order linear nonhomogeneous ODE (constant coefficients)<br/></div>

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

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

        & \\displaystyle \\Rightarrow \\text{{Solve the homogeneous part: }} {L(a_)}y'' + {L(b_)}y' + {L(c_)}y = 0 \\\\
        & \\displaystyle \\Rightarrow \\text{{Characteristic equation: }} {char_ltx} \\\\
        & \\displaystyle \\Rightarrow {D_ltx} \\\\
        & \\displaystyle \\Rightarrow {r1_ltx} \\\\
        & \\displaystyle \\Rightarrow {r2_ltx} \\\\
        & \\displaystyle \\Rightarrow {yc_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Find a particular solution }} y_p \\text{{ for }} g(x) \\\\
        & \\displaystyle \\Rightarrow \\text{{{method_note}}} \\\\
        & \\displaystyle \\Rightarrow {yp_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{General solution:}} \\\\
        & \\displaystyle \\Rightarrow {ygen_ltx} \\\\

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

    # IC block
    ic_ltx = ""
    if use_ic and (y_particular is not None) and (C1_value is not None) and (C2_value is not None):
        x0_ = sp.sympify(x0, locals=base_locals)
        y0_ = sp.sympify(y0, locals=base_locals)
        dy0_ = sp.sympify(dy0, locals=base_locals)

        yprime_general = sp.diff(y_general, x)

        ic1_ltx = f"{L(yx)}\\Big|_{{{L(x)}={L(x0_)}}} = {L(y0_)}"
        ic2_ltx = f"{L(y1)}\\Big|_{{{L(x)}={L(x0_)}}} = {L(dy0_)}"

        # Explicit “at x=x0” substitution lines (your style)
        y_at_x0_ltx = f"{L(y_general.subs(x, x0_))} = {L(y0_)} \\qquad \\text{{at }} {L(x)}={L(x0_)}"
        yp_at_x0_ltx = f"{L(yprime_general.subs(x, x0_))} = {L(dy0_)} \\qquad \\text{{at }} {L(x)}={L(x0_)}"

        C1_ltx = f"{L(C1)} = {L(C1_value)}"
        C2_ltx = f"{L(C2)} = {L(C2_value)}"

        final_ic_soln_ltx = f"{L(yx)} = {L(y_particular)} \\qquad \\cdots \\text{{from (iii)}}"
        final_ic_ans_ltx = f"{L(yx)} = {L(y_particular)}"

        ic_block = f"""
            <br/>
            \\(\\displaystyle
            \\begin{{aligned}}
            & \\displaystyle \\Rightarrow \\text{{Given }} {ic1_ltx},\\; {ic2_ltx} \\\\
            & \\displaystyle \\Rightarrow {y_at_x0_ltx} \\\\
            & \\displaystyle \\Rightarrow {yp_at_x0_ltx} \\\\
            & \\displaystyle \\Rightarrow {C1_ltx},\\; {C2_ltx} \\\\
            & \\displaystyle \\Rightarrow {final_ic_soln_ltx} \\\\
            \\end{{aligned}}
            \\)
        """
        soln_stmts.append(ic_block.replace("\n", ""))

        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {ic1_ltx},\\; {ic2_ltx} \\)"
        answer_s = f"\\( {final_ic_ans_ltx} \\)"
    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(),
    }


# ============================================================
# Optional: KProblemTemplate-style generator method (like yours)
# ============================================================

