# ============================================================
# File: _solve_linear_systems_nxn_matrix_exponential.py
# Topic: General (n×n) Linear Systems
#
#   x'(t) = A x(t)                      (homogeneous)
#   x'(t) = A x(t) + b(t)               (nonhomogeneous)
#
# Core facts:
#   - Fundamental matrix Φ(t) solves Φ' = AΦ, Φ(0)=I, hence Φ(t)=exp(A t)
#   - Homogeneous solution:  x(t) = exp(A t) x(0)
#   - Nonhomogeneous (variation of parameters / convolution):
#         x(t) = exp(A(t-t0)) x(t0) + ∫_{t0}^{t} exp(A(t-s)) b(s) ds
#
# This solver is "app-friendly":
#   - Takes A (list-of-lists or SymPy Matrix or string) and b(t) optional
#   - Computes exp(A t) using SymPy's expm (MatrixExp) and simplifies
#   - Reports eigenvalues (spectrum) and stability (Re(λ)<0) if decidable
#   - Optionally substitutes IC x(t0)=x0
#   - For b(t):
#       * if constant vector, SymPy often integrates exp(A(t-s)) b ds
#       * otherwise returns integral form (still correct)
#
# Notes:
#   - Avoid heavy Jordan/diagonalization code; SymPy handles expm generally.
#   - We still show diagonalizable/Jordan explanation in steps (text),
#     and we present exp(A t) as the fundamental matrix.
# ============================================================

from __future__ import annotations

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

import sympy as sp


def solve_linear_systems_nxn_matrix_exponential(
    A_in: Union[str, sp.Matrix, Sequence[Sequence[Union[int, float, str, sp.Expr]]]],
    *,
    b_in: Optional[Union[str, Sequence[Union[int, float, str, sp.Expr]], sp.Matrix]] = None,
    # symbols
    t_symbol: str = "t",
    x_symbol: str = "x",
    t0: Union[int, float, str, sp.Expr] = 0,
    x0: Optional[Sequence[Union[int, float, str, sp.Expr]]] = None,
    simplify: bool = True,
    locals_map: Optional[Dict[str, Any]] = None,
) -> dict:
    """
    Solve x'(t)=A x(t) (+ b(t)) using matrix exponential / fundamental matrix.

    Returns:
      { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
    """
    t = sp.Symbol(t_symbol, real=True)

    base_locals = {
        t_symbol: t,
        "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,
        "sinh": sp.sinh,
        "cosh": sp.cosh,
        "tanh": sp.tanh,
        "Abs": sp.Abs,
        "Matrix": sp.Matrix,
    }
    if locals_map:
        base_locals.update(locals_map)

    def P(v: Union[int, float, str, sp.Expr]) -> sp.Expr:
        return sp.sympify(v, locals=base_locals)

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

    # -----------------------------
    # Parse A
    # -----------------------------
    if isinstance(A_in, sp.MatrixBase):
        A = sp.Matrix(A_in)
    elif isinstance(A_in, str):
        try:
            A_expr = sp.sympify(A_in, locals=base_locals)
            if isinstance(A_expr, sp.MatrixBase):
                A = sp.Matrix(A_expr)
            else:
                raise ValueError("A_in string must parse to a SymPy Matrix.")
        except Exception as e:
            raise ValueError(f"Could not parse A_in={A_in!r}. Error: {e}") from e
    else:
        # list-of-lists
        try:
            A = sp.Matrix([[P(v) for v in row] for row in A_in])
        except Exception as e:
            raise ValueError(f"Could not build matrix from A_in. Error: {e}") from e

    n = int(A.rows)
    if A.cols != n:
        raise ValueError("A must be a square matrix (n×n).")

    if simplify:
        A = sp.simplify(A)

    # -----------------------------
    # Parse b(t) (optional)
    # -----------------------------
    b = None
    if b_in is not None:
        if isinstance(b_in, sp.MatrixBase):
            b = sp.Matrix(b_in)
        elif isinstance(b_in, str):
            # allow b as a vector string like "[sin(t), 0, exp(t)]"
            try:
                b_expr = sp.sympify(b_in, locals=base_locals)
                if isinstance(b_expr, sp.MatrixBase):
                    b = sp.Matrix(b_expr)
                elif isinstance(b_expr, (list, tuple)):
                    b = sp.Matrix([sp.sympify(v, locals=base_locals) for v in b_expr])
                else:
                    # scalar b not allowed
                    raise ValueError("b_in string must parse to a vector (list/tuple/Matrix).")
            except Exception as e:
                raise ValueError(f"Could not parse b_in={b_in!r}. Error: {e}") from e
        else:
            # sequence
            try:
                b = sp.Matrix([P(v) for v in b_in])
            except Exception as e:
                raise ValueError(f"Could not parse b_in sequence. Error: {e}") from e

        # ensure shape (n,1)
        b = sp.Matrix(b)
        if b.rows == 1 and b.cols == n:
            b = b.T
        if b.cols != 1 or b.rows != n:
            raise ValueError(f"b(t) must be a vector of length {n}. Got shape {b.shape}.")

        if simplify:
            b = sp.simplify(b)

    # -----------------------------
    # Initial condition
    # -----------------------------
    t0e = P(t0)
    if simplify:
        t0e = sp.simplify(t0e)

    x0_vec = None
    use_ic = x0 is not None
    if use_ic:
        if len(x0) != n:
            raise ValueError(f"x0 must have length {n}.")
        x0_vec = sp.Matrix([P(v) for v in x0])
        if simplify:
            x0_vec = sp.simplify(x0_vec)

    # -----------------------------
    # Fundamental matrix Φ(t) = exp(A (t-t0))
    # -----------------------------
    tau = sp.Symbol("tau", real=True)
    Phi = sp.exp(A * (t - t0e))
    if simplify:
        Phi = sp.simplify(Phi)

    # -----------------------------
    # Build solution
    # -----------------------------
    x_t = None
    integral_part = None

    if b is None:
        # Homogeneous
        if use_ic:
            x_t = Phi * x0_vec
        else:
            c = sp.Matrix(sp.symbols(f"c1:{n+1}")).reshape(n, 1)
            x_t = Phi * c

    else:
        # Nonhomogeneous: x(t) = Phi(t) x0 + ∫_{t0}^{t} exp(A(t-s)) b(s) ds
        s = sp.Symbol("s", real=True)
        kernel = sp.exp(A * (t - s))
        if simplify:
            kernel = sp.simplify(kernel)

        b_s = b.subs(t, s)
        integrand = kernel * b_s
        integral_part = sp.Integral(integrand, (s, t0e, t))

        if use_ic:
            x_t = Phi * x0_vec + integral_part
        else:
            c = sp.Matrix(sp.symbols(f"c1:{n+1}")).reshape(n, 1)
            x_t = Phi * c + integral_part

    if simplify:
        x_t = sp.simplify(x_t)

    # -----------------------------
    # Spectrum / stability hint
    # -----------------------------
    eigvals = []
    stable_hint = "Stability: inconclusive (symbolic eigenvalues)."
    try:
        ev_dict = A.eigenvals()
        eigvals = list(ev_dict.keys())

        # check if all Re(ev) < 0 when possible
        re_list = []
        all_decidable = True
        all_neg = True
        for ev in eigvals:
            re_ev = sp.re(ev)
            re_list.append(re_ev)
            sgn = sp.sign(re_ev)
            if sgn not in (sp.Integer(-1), -1, sp.Integer(0), 0, sp.Integer(1), 1):
                all_decidable = False
            else:
                # treat 0 as not asymptotically stable
                if sgn in (sp.Integer(0), 0) or sgn in (sp.Integer(1), 1):
                    all_neg = False

        if all_decidable:
            stable_hint = "Stability: asymptotically stable (all Re(λ)<0)." if all_neg else "Stability: not asymptotically stable (some Re(λ)≥0)."
        else:
            stable_hint = "Stability: based on eigenvalues; asymptotic stability requires all Re(λ)<0."
    except Exception:
        eigvals = []
        stable_hint = "Stability: eigenvalue computation failed."

    # -----------------------------
    # Build HTML output (no f-string brace hazards)
    # -----------------------------
    A_ltx = Ltx(A)
    Phi_ltx = Ltx(Phi)

    if b is None:
        sys_ltx = "\\mathbf{x}'(t)=A\\mathbf{x}(t),\\quad A=" + A_ltx
    else:
        sys_ltx = "\\mathbf{x}'(t)=A\\mathbf{x}(t)+\\mathbf{b}(t),\\quad A=" + A_ltx + ",\\; \\mathbf{b}(t)=" + Ltx(b)

    problem_s = "Solve the linear system: \\( " + sys_ltx + " \\)."

    eig_ltx = ""
    if eigvals:
        eig_ltx = ",\\; ".join([Ltx(ev) for ev in eigvals])

    # x(t) latex
    xvec = sp.MatrixSymbol(x_symbol, n, 1)
    # display x(t) as explicit vector expression
    x_t_ltx = Ltx(x_t)

    steps = (
        "& \\displaystyle \\Rightarrow \\text{Fundamental matrix }\\Phi(t)\\text{ solves }\\Phi'=A\\Phi,\\;\\Phi(t_0)=I \\\\"
        "& \\displaystyle \\Rightarrow \\Phi(t)=e^{A(t-t_0)} \\\\"
        "& \\displaystyle \\Rightarrow \\Phi(t)= " + Phi_ltx + " \\\\[10pt]"
    )

    if b is None:
        if use_ic:
            steps += (
                "& \\displaystyle \\Rightarrow \\mathbf{x}(t)=\\Phi(t)\\,\\mathbf{x}(t_0) \\\\"
                "& \\displaystyle \\Rightarrow \\mathbf{x}(t)= " + x_t_ltx + " \\\\"
            )
        else:
            steps += (
                "& \\displaystyle \\Rightarrow \\mathbf{x}(t)=\\Phi(t)\\,\\mathbf{c} \\text{ (constant vector)} \\\\"
                "& \\displaystyle \\Rightarrow \\mathbf{x}(t)= " + x_t_ltx + " \\\\"
            )
    else:
        steps += (
            "& \\displaystyle \\Rightarrow \\text{Variation of parameters (matrix form):} \\\\"
            "& \\displaystyle \\Rightarrow \\mathbf{x}(t)=\\Phi(t)\\mathbf{x}(t_0)+\\int_{t_0}^{t} e^{A(t-s)}\\mathbf{b}(s)\\,ds \\\\"
            "& \\displaystyle \\Rightarrow \\mathbf{x}(t)= " + x_t_ltx + " \\\\"
        )

    spec_block = ""
    if eigvals:
        spec_block = (
            "& \\displaystyle \\Rightarrow \\text{Eigenvalues of }A:\\; " + eig_ltx + " \\\\"
            "& \\displaystyle \\Rightarrow \\text{" + stable_hint + "} \\\\"
        )
    else:
        spec_block = "& \\displaystyle \\Rightarrow \\text{" + stable_hint + "} \\\\"

    soln_s = (
        "<div class='wrap'><div class='card'>"
        "<div class='badge'>General Linear Systems (n×n) — Matrix Exponential & Fundamental Matrix<br/></div>"
        "\\(\\displaystyle \\begin{aligned}"
        "& \\displaystyle " + sys_ltx + " \\\\[10pt]"
        + steps +
        spec_block +
        "\\end{aligned}\\)"
        "<div class='note'>If \\(A\\) is diagonalizable, \\(e^{At}=P e^{Dt}P^{-1}\\). If not, use Jordan blocks: \\(e^{Jt}=e^{\\lambda t}(I+Nt+\\cdots)\\).</div>"
        "</div></div>"
    )

    # Answer (compact)
    if b is None:
        if use_ic:
            answer_s = "\\( \\mathbf{x}(t)=e^{A(t-t_0)}\\mathbf{x}(t_0) = " + x_t_ltx + " \\)"
        else:
            answer_s = "\\( \\mathbf{x}(t)=e^{A(t-t_0)}\\mathbf{c} = " + x_t_ltx + " \\)"
    else:
        answer_s = "\\( \\mathbf{x}(t)=e^{A(t-t_0)}\\mathbf{x}(t_0)+\\int_{t_0}^{t} e^{A(t-s)}\\mathbf{b}(s)ds = " + x_t_ltx + " \\)"

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