# ============================================================
# File: _solve_ode_linear_system_2x2_constant_coeff.py
# Topic: Linear System of ODEs (2x2, constant coefficients)
#
# Form:
#   X'(x) = A X(x)
# where X(x) = [u(x), v(x)]^T and A is a constant 2x2 matrix.
#
# Solution:
#   X(x) = exp(A (x-x0)) X(x0)          (IVP form)
#   X(x) = exp(A x) C                   (general form)
#
# Notes:
#   - Uses SymPy matrix exponential (A*t).exp()
#   - Shows eigenvalues (best-effort) for classification.
# ============================================================

from __future__ import annotations

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

import sympy as sp


def solve_ode_linear_system_2x2_constant_coeff(
    A_2x2: Union[List[List[Union[str, int, float, sp.Expr]]], sp.Matrix],
    *,
    x_symbol: str = "x",
    u_symbol: str = "u",
    v_symbol: str = "v",
    C1_symbol: str = "C_1",
    C2_symbol: str = "C_2",
    x0: Optional[Union[int, float, str, sp.Expr]] = None,
    u0: Optional[Union[int, float, str, sp.Expr]] = None,
    v0: Optional[Union[int, float, str, sp.Expr]] = None,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Solve X'(x) = A X(x), where X=[u,v]^T, A is 2x2 constant matrix.

    Optional IC at x0: u(x0)=u0, v(x0)=v0

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

    # -----------------------------
    # 1) Symbols + parsing
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    t = sp.Symbol("t", real=True)  # internal shift variable

    u = sp.Function(u_symbol)
    v = sp.Function(v_symbol)
    ux = u(x)
    vx = v(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,
        "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(vv):
        if isinstance(vv, str):
            return sp.sympify(vv, locals=base_locals)
        return sp.sympify(vv)

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

    # Matrix A
    if isinstance(A_2x2, sp.Matrix):
        A = A_2x2
    else:
        A = sp.Matrix([[ _parse(A_2x2[0][0]), _parse(A_2x2[0][1]) ],
                       [ _parse(A_2x2[1][0]), _parse(A_2x2[1][1]) ]])

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

    # IC
    use_ic = (x0 is not None) and (u0 is not None) and (v0 is not None)
    X0 = None
    parsed_x0 = None
    if use_ic:
        parsed_x0 = _parse(x0)
        X0 = sp.Matrix([_parse(u0), _parse(v0)])
        if simplify:
            X0 = sp.Matrix([sp.simplify(X0[0]), sp.simplify(X0[1])])

    # -----------------------------
    # 2) Compute matrix exponential
    # -----------------------------
    # General: X(x) = exp(A*(x)) * C
    # IVP:     X(x) = exp(A*(x-x0)) * X0
    try:
        expA_x = (A * x).exp()
    except Exception:
        expA_x = sp.exp(A * x)  # fallback symbolic

    # Shifted exp
    if use_ic:
        dx = sp.simplify(x - parsed_x0) if simplify else (x - parsed_x0)
        try:
            expA_dx = (A * dx).exp()
        except Exception:
            expA_dx = sp.exp(A * dx)
        X_expr = expA_dx * X0
        if simplify:
            X_expr = sp.Matrix([sp.simplify(X_expr[0]), sp.simplify(X_expr[1])])
    else:
        Cvec = sp.Matrix([C1, C2])
        X_expr = expA_x * Cvec
        if simplify:
            X_expr = sp.Matrix([sp.simplify(X_expr[0]), sp.simplify(X_expr[1])])

    # Eigenvalues (best-effort)
    evals = None
    try:
        evals = A.eigenvals()  # {eigenvalue: mult}
    except Exception:
        evals = None

    # -----------------------------
    # 3) Build LaTeX / HTML
    # -----------------------------
    # System form:
    ueq = sp.Eq(sp.diff(ux, x), A[0, 0] * ux + A[0, 1] * vx)
    veq = sp.Eq(sp.diff(vx, x), A[1, 0] * ux + A[1, 1] * vx)

    ueq_ltx = f"\\frac{{d{L(ux)}}}{{d{L(x)}}} = {L(A[0,0])}{L(ux)} + {L(A[0,1])}{L(vx)}"
    veq_ltx = f"\\frac{{d{L(vx)}}}{{d{L(x)}}} = {L(A[1,0])}{L(ux)} + {L(A[1,1])}{L(vx)}"

    A_ltx = L(A)
    X_ltx = f"\\mathbf{{X}}({L(x)})=\\begin{{pmatrix}}{L(ux)}\\\\{L(vx)}\\end{{pmatrix}}"

    if use_ic:
        ic_ltx = f"{L(u)}({L(parsed_x0)})={L(X0[0])},\\; {L(v)}({L(parsed_x0)})={L(X0[1])}"
        exp_ltx = f"\\mathbf{{X}}(x)=\\exp\\left(A(x-{L(parsed_x0)})\\right)\\,\\mathbf{{X}}({L(parsed_x0)})"
    else:
        ic_ltx = ""
        exp_ltx = f"\\mathbf{{X}}(x)=\\exp\\left(Ax\\right)\\,\\mathbf{{C}},\\; \\mathbf{{C}}=\\begin{{pmatrix}}{L(C1)}\\\\{L(C2)}\\end{{pmatrix}}"

    # Eigenvalues display
    eig_ltx = ""
    note_parts = []
    if evals:
        eig_lines = []
        for lam, mult in evals.items():
            eig_lines.append(f"& \\displaystyle \\lambda={L(lam)}\\;\\; (\\text{{mult}}={mult})")
        eig_ltx = " \\\\ ".join(eig_lines)
    else:
        note_parts.append("Eigenvalues could not be extracted cleanly; matrix exponential is used directly.")

    # Final solution components
    u_sol = X_expr[0]
    v_sol = X_expr[1]

    anti_note = " ".join(note_parts)

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: Linear System (2×2) with Constant Coefficients<br/></div>

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

        & \\displaystyle {ueq_ltx} \\\\
        & \\displaystyle {veq_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Write in vector form: }} \\mathbf{{X}}' = A\\mathbf{{X}} \\\\
        & \\displaystyle \\Rightarrow {X_ltx},\\quad A={A_ltx} \\\\[10pt]

        & \\displaystyle \\Rightarrow \\text{{Solution via matrix exponential:}} \\\\
        & \\displaystyle \\Rightarrow {exp_ltx} \\\\[10pt]
    """

    if eig_ltx:
        soln_stmt += f"""
        & \\displaystyle \\Rightarrow \\text{{Eigenvalues of }}A: \\\\
        {eig_ltx} \\\\[10pt]
        """

    soln_stmt += f"""
        & \\displaystyle \\Rightarrow \\begin{{pmatrix}}{L(ux)}\\\\{L(vx)}\\end{{pmatrix}}
          = \\begin{{pmatrix}}{L(u_sol)}\\\\{L(v_sol)}\\end{{pmatrix}} \\\\

        \\end{{aligned}}
        \\)

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

    if use_ic:
        problem_s = f"Solve the system using matrix exponential with {ic_ltx}: \\( {ueq_ltx},\\; {veq_ltx} \\)"
        answer_s = f"\\( {L(ux)}={L(u_sol)},\\quad {L(vx)}={L(v_sol)} \\)"
    else:
        problem_s = f"Solve the system using matrix exponential: \\( {ueq_ltx},\\; {veq_ltx} \\)"
        answer_s = f"\\( {L(ux)}={L(u_sol)},\\quad {L(vx)}={L(v_sol)} \\)"

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

