from __future__ import annotations

from typing import Optional, Union, Dict, Any

import sympy as sp


def solve_ode_second_order_linear_homogeneous_const_coeff(
    a: Union[str, sp.Expr],
    b: Union[str, sp.Expr],
    c: 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:
    """
    Generates an HTML "full steps" solution for the ODE:
        a y'' + b y' + c y = 0
    where a, b, c are constants (with respect to x) and a != 0.

    Optional initial conditions (provide all three for IC mode):
        y(x0) = y0
        y'(x0) = dy0

    Output:
      dict with keys: problem_s, soln_s, answer_s
    """

    # -----------------------------
    # 1) Setup symbols + parsing
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    yx = sp.Function(y_symbol)(x)
    r = sp.Symbol(r_symbol)
    C1 = sp.Symbol(C1_symbol)
    C2 = sp.Symbol(C2_symbol)

    y1 = sp.diff(yx, x)
    y2 = sp.diff(yx, x, 2)

    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,
    }

    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)
    except Exception as e:
        raise ValueError(f"Could not parse (a,b,c)=({a!r},{b!r},{c!r}). Error: {e}") from e

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

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

    # -----------------------------
    # 2) Characteristic equation
    # -----------------------------
    # Trial y = e^{rx}
    y_trial = sp.exp(r * x)

    # NOTE: We keep the characteristic polynomial explicit for the written steps.
    char_poly = a_ * r**2 + b_ * r + c_
    D = sp.simplify(b_**2 - 4 * a_ * c_) if simplify else (b_**2 - 4 * a_ * c_)

    # Quadratic roots (symbolic)
    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_))

    # Determine case if SymPy can decide
    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"

    # Build general solution expression
    if case == "repeated":
        r0 = r1
        y_general_expr = (C1 + C2 * x) * sp.exp(r0 * x)
        solution_form_note = r"D=0 \Rightarrow \text{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_))
        y_general_expr = sp.exp(alpha * x) * (C1 * sp.cos(beta * x) + C2 * sp.sin(beta * x))
        solution_form_note = r"D<0 \Rightarrow \text{complex roots } \alpha \pm i\beta"
    else:
        # distinct real OR unknown -> use exp(r1 x), exp(r2 x) form
        y_general_expr = C1 * sp.exp(r1 * x) + C2 * sp.exp(r2 * x)
        solution_form_note = (
            r"D>0 \Rightarrow \text{distinct real roots}"
            if case == "distinct_real"
            else r"\text{Use the roots } r_1, r_2"
        )

    if simplify:
        y_general_expr = sp.simplify(y_general_expr)

    # -----------------------------
    # 3) Optional initial conditions
    # -----------------------------
    use_ic = (x0 is not None) and (y0 is not None) and (dy0 is not None)
    C1_value = None
    C2_value = None
    y_particular_expr = 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_expr = sp.diff(y_general_expr, x)

        eq1 = sp.Eq(y_general_expr.subs(x, x0_), y0_)
        eq2 = sp.Eq(yprime_expr.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_expr = y_general_expr.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_expr = sp.simplify(y_particular_expr)

    # -----------------------------
    # 4) Build solution HTML (similar style)
    # -----------------------------
    ode_ltx = rf"{L(a_)}\,{L(y2)} + {L(b_)}\,{L(y1)} + {L(c_)}\,{L(yx)} = 0"

    y_trial_ltx = rf"{L(yx)} = e^{{{L(r)}{L(x)}}}"
    y_trial_1_ltx = rf"{L(y1)} = {L(r)}\,e^{{{L(r)}{L(x)}}}"
    y_trial_2_ltx = rf"{L(y2)} = {L(r)}^2\,e^{{{L(r)}{L(x)}}}"

    subs_ltx = rf"{L(a_)}({L(r)}^2 e^{{{L(r)}{L(x)}}}) + {L(b_)}({L(r)} e^{{{L(r)}{L(x)}}}) + {L(c_)}(e^{{{L(r)}{L(x)}}}) = 0"
    fact_ltx = rf"\left({L(a_)}{L(r)}^2 + {L(b_)}{L(r)} + {L(c_)}\right)e^{{{L(r)}{L(x)}}} = 0"
    char_ltx = rf"{L(a_)}{L(r)}^2 + {L(b_)}{L(r)} + {L(c_)} = 0 \qquad \cdots \text{{(Characteristic)}}"

    D_ltx = rf"D = {L(b_)}^2 - 4({L(a_)})({L(c_)}) = {L(D)}"
    r1_ltx = rf"{L(r)}_1 = \frac{{-{L(b_)} + \sqrt{{{L(D)}}}}}{{2({L(a_)})}} = {L(r1)}"
    r2_ltx = rf"{L(r)}_2 = \frac{{-{L(b_)} - \sqrt{{{L(D)}}}}}{{2({L(a_)})}} = {L(r2)}"

    general_ltx = rf"{L(yx)} = {L(y_general_expr)} \qquad \cdots \text{{(i)}}"

    soln_stmts = []

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

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

        & \\displaystyle {ode_ltx} \\\\[6pt]
        & \\displaystyle \\text{{Assume a trial solution: }} {y_trial_ltx} \\\\
        & \\displaystyle \\implies {y_trial_1_ltx},\\quad {y_trial_2_ltx} \\\\[6pt]
        & \\displaystyle \\text{{Substitute into the ODE:}} \\\\
        & \\displaystyle \\implies {subs_ltx} \\\\
        & \\displaystyle \\implies {fact_ltx} \\\\
        & \\displaystyle \\implies {char_ltx} \\\\[6pt]
        & \\displaystyle {D_ltx} \\\\
        & \\displaystyle \\implies {r1_ltx} \\\\
        & \\displaystyle \\implies {r2_ltx} \\\\[6pt]
        & \\displaystyle \\text{{Therefore, }} {solution_form_note} \\\\
        & \\displaystyle \\implies {general_ltx} \\\\

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

    ic_block = ""
    if use_ic and (C1_value is not None) and (C2_value is not None) and (y_particular_expr 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_expr = sp.diff(y_general_expr, x)

        # Extra "substitute x=x0" style lines (like the first-order solver screenshot)
        yprime_general_ltx = rf"{L(y1)} = {L(yprime_expr)} \qquad \cdots \text{{(ii)}}"
        y_at_x0_from_i_ltx = rf"{L(yx.subs(x, x0_))} = {L(y_general_expr.subs(x, x0_))} \qquad \text{{at }} {L(x)}={L(x0_)}"
        yp_at_x0_from_ii_ltx = rf"{L(y1.subs(x, x0_))} = {L(yprime_expr.subs(x, x0_))} \qquad \text{{at }} {L(x)}={L(x0_)}"

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

        y_at_x0_ltx = rf"{L(y_general_expr.subs(x, x0_))} = {L(y0_)}"
        yp_at_x0_ltx = rf"{L(yprime_expr.subs(x, x0_))} = {L(dy0_)}"

        c1_ltx = rf"{L(C1)} = {L(C1_value)} \qquad \text{{at }} {L(x)}={L(x0_)}"
        c2_ltx = rf"{L(C2)} = {L(C2_value)} \qquad \text{{at }} {L(x)}={L(x0_)}"

        final_ic_ltx = rf"{L(yx)} = {L(y_particular_expr)}"

        ic_block = f"""
            <br/>
            \\(\\displaystyle
            \\begin{{aligned}}

            & \\displaystyle \\text{{Given }} {ic1_ltx},\\; {ic2_ltx} \\\\
            & \\displaystyle \\implies {yprime_general_ltx} \\\\
            & \\displaystyle \\implies {y_at_x0_from_i_ltx} \\\\
            & \\displaystyle \\implies {yp_at_x0_from_ii_ltx} \\\\
            & \\displaystyle \\implies {y_at_x0_ltx} \\\\
            & \\displaystyle \\implies {yp_at_x0_ltx} \\\\
            & \\displaystyle \\implies {c1_ltx} \\\\
            & \\displaystyle \\implies {c2_ltx} \\\\
            & \\displaystyle \\implies {final_ic_ltx}  \\qquad \\cdots \\text{{ from (i)}} \\\\

            \\end{{aligned}}
            \\)
        """
        soln_stmts.append(ic_block.replace("\n", ""))

    # -----------------------------
    # 5) Final problem/answer strings
    # -----------------------------
    if use_ic and ic_block:
        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {ic1_ltx},\\; {ic2_ltx} \\)"
        answer_s = f"\\( {L(yx)} = {L(y_particular_expr)} \\)"
    else:
        problem_s = f"Solve \\( {ode_ltx} \\)"
        answer_s = f"\\( {L(yx)} = {L(y_general_expr)} \\)"

    soln_s = "".join(soln_stmts)

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