# ============================================================
# File: _solve_ode_laplace_transform_ivp.py
# Topic: Laplace Transform Method for IVPs (Linear ODE, constant coefficients)
#
# Form:
#   a_n y^(n) + a_{n-1} y^(n-1) + ... + a_1 y' + a_0 y = g(x)
# with initial conditions at x=0:
#   y(0)=y0, y'(0)=y1, ..., y^(n-1)(0)=y_{n-1}
#
# Method:
#   1) Take Laplace transform of both sides.
#   2) Use derivative transforms:
#        L{y'} = sY - y(0)
#        L{y''} = s^2 Y - s y(0) - y'(0)
#        ...
#   3) Solve algebraic equation for Y(s).
#   4) Inverse Laplace to get y(x).
#
# Notes:
#   - If SymPy can’t invert, it keeps inverse_laplace_transform(...)
#   - Works well for polynomials/exp/sin/cos/Heaviside/DiracDelta (best-effort)
# ============================================================

from __future__ import annotations

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

import sympy as sp


def solve_ode_laplace_transform_ivp(
    coeffs: List[Union[str, int, float, sp.Expr]],
    rhs_of_x: Union[str, sp.Expr],
    ics: Dict[int, Union[str, int, float, sp.Expr]],
    *,
    x_symbol: str = "x",
    y_symbol: str = "y",
    s_symbol: str = "s",
    Y_symbol: str = "Y",
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Solve linear IVP by Laplace transform.

    coeffs: [a_n, a_{n-1}, ..., a_0] (constants), length n+1
    rhs_of_x: g(x)
    ics: dict {k: y^(k)(0)} for k=0..n-1 (missing treated as 0)

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

    # -----------------------------
    # 1) Symbols + parsing
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    s = sp.Symbol(s_symbol, real=True, positive=True)

    y = sp.Function(y_symbol)
    yx = y(x)

    # Laplace domain unknown
    Y = sp.Function(Y_symbol)
    Ys = sp.Symbol("Y")  # we will use this symbol for algebra (Y(s))

    base_locals = {
        x_symbol: x,
        s_symbol: s,
        "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,
        "Heaviside": sp.Heaviside,
        "DiracDelta": sp.DiracDelta,
        "Piecewise": sp.Piecewise,
        "sec": sp.sec,
        "csc": sp.csc,
        "cot": sp.cot,
    }
    if locals_map:
        base_locals.update(locals_map)

    def _parse(v):
        if isinstance(v, str):
            return sp.sympify(v, locals=base_locals)
        return sp.sympify(v)

    def L(expr: sp.Expr) -> str:
        return sp.latex(expr)

    if not coeffs or len(coeffs) < 2:
        raise ValueError("coeffs must be [a_n, ..., a_0] with length >= 2.")

    a_list = [_parse(c) for c in coeffs]
    if simplify:
        a_list = [sp.simplify(a) for a in a_list]

    n = len(a_list) - 1
    if simplify:
        if sp.simplify(a_list[0]) == 0:
            raise ValueError("Leading coefficient a_n must be nonzero.")
    else:
        if a_list[0] == 0:
            raise ValueError("Leading coefficient a_n must be nonzero.")

    rhs = _parse(rhs_of_x)
    rhs = sp.simplify(rhs) if simplify else rhs

    # Parse ICs (default 0 if missing)
    ic_vals = {}
    for k in range(n):
        v = ics.get(k, 0)
        ic_vals[k] = sp.simplify(_parse(v)) if simplify else _parse(v)

    # -----------------------------
    # 2) Build Laplace transform of LHS using derivative formulas
    # -----------------------------
    # L{y^(m)} = s^m Y(s) - s^{m-1} y(0) - s^{m-2} y'(0) - ... - y^(m-1)(0)
    def laplace_derivative(m: int) -> sp.Expr:
        expr = (s**m) * Ys
        for j in range(m):
            expr -= (s**(m - 1 - j)) * ic_vals.get(j, 0)
        return sp.simplify(expr) if simplify else expr

    # LHS in s:
    lhs_s = sp.Integer(0)
    for idx, a_k in enumerate(a_list):
        order = n - idx
        if order == 0:
            lhs_s += a_k * Ys
        else:
            lhs_s += a_k * laplace_derivative(order)

    lhs_s = sp.simplify(lhs_s) if simplify else lhs_s

    # RHS Laplace
    try:
        rhs_s = sp.laplace_transform(rhs, x, s, noconds=True)
        rhs_s = sp.simplify(rhs_s) if simplify else rhs_s
    except Exception:
        rhs_s = sp.laplace_transform(rhs, x, s)[0]  # fallback
        rhs_s = sp.simplify(rhs_s) if simplify else rhs_s

    # Solve for Y(s)
    eq = sp.Eq(lhs_s, rhs_s)
    Ys_expr = None
    try:
        sol = sp.solve(eq, Ys)
        if sol:
            Ys_expr = sp.simplify(sol[0]) if simplify else sol[0]
    except Exception:
        Ys_expr = None

    if Ys_expr is None:
        # graceful fallback
        ode_ltx, ic_ltx = _latex_ode_and_ics(a_list, rhs, ic_vals, x, yx, L)
        problem_s = f"Solve the IVP using Laplace transforms: \\( {ode_ltx} \\), {ic_ltx}"
        soln_s = (
            "<div class='wrap'><div class='card'>"
            "<div class='badge'>Note: Laplace Transform Method (IVP)<br/></div>"
            "\\(\\displaystyle \\begin{aligned}"
            f"& \\displaystyle \\text{{Laplace equation: }} {L(lhs_s)} = {L(rhs_s)} \\\\"
            "& \\displaystyle \\Rightarrow \\text{Could not solve algebraically for }Y(s). \\\\"
            "\\end{aligned}\\)"
            "</div></div>"
        )
        answer_s = "\\(\\text{Could not solve for }Y(s).\\)"
        return {"problem_s": problem_s.strip(), "soln_s": soln_s.strip(), "answer_s": answer_s.strip()}

    # Inverse Laplace
    try:
        y_expr = sp.inverse_laplace_transform(Ys_expr, s, x)
        y_expr = sp.simplify(y_expr) if simplify else y_expr
    except Exception:
        y_expr = sp.inverse_laplace_transform(Ys_expr, s, x)  # keep symbolic

    # -----------------------------
    # 3) HTML formatting (match your style)
    # -----------------------------
    ode_ltx, ic_ltx = _latex_ode_and_ics(a_list, rhs, ic_vals, x, yx, L)

    # Build a compact derivative-transform display
    # show up to y'', y''', ... depending on n (cap to 4 lines)
    deriv_lines = []
    max_show = min(n, 4)
    for m in range(1, max_show + 1):
        if m == 1:
            deriv_lines.append(
                f"& \\displaystyle \\mathcal{{L}}\\{{y'\\}} = sY - y(0) = {L(laplace_derivative(1))}"
            )
        elif m == 2:
            deriv_lines.append(
                f"& \\displaystyle \\mathcal{{L}}\\{{y''\\}} = s^2Y - s y(0) - y'(0) = {L(laplace_derivative(2))}"
            )
        else:
            deriv_lines.append(
                f"& \\displaystyle \\mathcal{{L}}\\{{y^{{({m})}}\\}} = {L(laplace_derivative(m))}"
            )
    deriv_block = " \\\\ ".join(deriv_lines)

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: Laplace Transform Method (IVP)<br/></div>

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

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

        & \\displaystyle \\Rightarrow \\text{{Take Laplace transforms. Let }} Y(s)=\\mathcal{{L}}\\{{y(x)\\}}. \\\\[6pt]
        {deriv_block} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Transform the whole ODE:}} \\\\
        & \\displaystyle {L(lhs_s)} = {L(rhs_s)} \\\\[8pt]

        & \\displaystyle \\Rightarrow Y(s) = {L(Ys_expr)} \\\\[10pt]

        & \\displaystyle \\Rightarrow y(x)=\\mathcal{{L}}^{{-1}}\\{{Y(s)\\}} = {L(y_expr)} \\\\

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

    problem_s = f"Solve the IVP using Laplace transforms: \\( {ode_ltx} \\), {ic_ltx}"
    answer_s = f"\\( {L(yx)} = {L(y_expr)} \\)"

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


def _latex_ode_and_ics(a_list: List[sp.Expr], rhs: sp.Expr, ic_vals: Dict[int, sp.Expr],
                      x: sp.Symbol, yx: sp.Expr, L) -> tuple:
    """
    Helper: Build latex for ODE and ICs at x=0.
    """
    n = len(a_list) - 1

    # ODE latex
    terms = []
    for idx, a_k in enumerate(a_list):
        order = n - idx
        if a_k == 0:
            continue
        if order == 0:
            t = f"\\left({L(a_k)}\\right){L(yx)}"
        elif order == 1:
            t = f"\\left({L(a_k)}\\right)\\frac{{d{L(yx)}}}{{d{L(x)}}}"
        else:
            t = f"\\left({L(a_k)}\\right)\\frac{{d^{order}{L(yx)}}}{{d{L(x)}^{order}}}"
        terms.append(t)

    ode_ltx = " + ".join(terms) + f" = {L(rhs)}"

    # IC latex
    ic_parts = []
    for k in range(n):
        val = ic_vals.get(k, 0)
        if k == 0:
            ic_parts.append(f"y(0)={L(val)}")
        elif k == 1:
            ic_parts.append(f"y'(0)={L(val)}")
        else:
            ic_parts.append(f"y^{{({k})}}(0)={L(val)}")

    ic_ltx = ",\\; ".join(ic_parts) if ic_parts else ""

    return ode_ltx, ic_ltx
