--- Final Solver (used for performance test) ---
import numpy as np
import scipy.linalg
from scipy.linalg._matfuncs import recursive_schur_sqrtm
import cmath
import warnings
from typing import Any

# ---------------------------------------------------------------------------
#  Suppress all scipy linalg warnings globally so individual solve() calls
#  don't pay the cost of Python's warning machinery processing messages.
# ---------------------------------------------------------------------------
warnings.filterwarnings("ignore", category=scipy.linalg.LinAlgWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning,
                        module=r"scipy\.linalg")
warnings.filterwarnings("ignore", category=FutureWarning,
                        module=r"scipy\.linalg")

# Pre-bind to avoid module attribute lookups in hot path
_scipy_sqrtm = scipy.linalg.sqrtm
_recursive_schur_sqrtm = recursive_schur_sqrtm


class Solver:
    """
    Fast matrix-square-root solver.

    Strategy (ordered by speed importance):
      n=1  → direct scalar complex sqrt          (≈13× faster than ref)
      n=2  → analytic closed-form (A + sI)/t     (≈6×  faster than ref)
      n=3,4→ compiled recursive_schur_sqrtm      (≈3×  faster than ref)
             called directly, bypassing the scipy wrapper overhead.
      n≥5  → scipy.linalg.sqrtm WITHOUT disp     (1.3-1.8× faster than ref)
             The reference uses disp=False which adds a full X@X multiply
             inside scipy for an error estimate that is never consumed.
             With ILP64 OpenBLAS this X@X multiply dominates for large n.

    All paths produce the same result as scipy's principal square root
    (eigenvalues with non-negative real parts) to within the tolerances
    used by is_solution (rtol=1e-5, atol=1e-8).
    """

    def solve(self, problem, **kwargs) -> Any:
        A = problem["matrix"]
        if not isinstance(A, np.ndarray):
            A = np.array(A, dtype=np.complex128)
        n = A.shape[0]

        # ---- n=1: scalar sqrt ----
        if n == 1:
            return {"sqrtm": {"X": [[cmath.sqrt(complex(A[0, 0]))]]}}

        # ---- n=2: analytic closed-form ----
        if n == 2:
            result = _sqrtm_2x2(A)
            if result is not None:
                return {"sqrtm": {"X": result}}

        # ---- n=3-4: bypass scipy wrapper, call compiled function directly ----
        if n <= 4:
            try:
                X, _, _, info = _recursive_schur_sqrtm(A)
                if info >= 0:
                    return {"sqrtm": {"X": X.tolist()}}
            except Exception:
                pass  # fall through to general path

        # ---- n≥5 (or fallback): scipy sqrtm WITHOUT disp ----
        # Omitting disp avoids the internal X@X error-estimate multiply
        # that disp=False triggers, saving 30-90% on large matrices.
        try:
            X = _scipy_sqrtm(A)
        except Exception:
            return {"sqrtm": {"X": []}}

        return {"sqrtm": {"X": X.tolist()}}


# ---------------------------------------------------------------------------
#  2×2 analytic principal square root
# ---------------------------------------------------------------------------
#  For A = [[a,b],[c,d]] with τ = tr(A), δ = det(A):
#    s = √δ,  t = √(τ + 2s),  X = (A + sI)/t
#  Branch of s is chosen to maximise |τ+2s| for numerical stability.
#  Cayley-Hamilton guarantees X² = A when t ≠ 0.
# ---------------------------------------------------------------------------
def _sqrtm_2x2(A):
    a = complex(A[0, 0])
    b = complex(A[0, 1])
    c = complex(A[1, 0])
    d = complex(A[1, 1])

    tau = a + d
    delta = a * d - b * c

    s1 = cmath.sqrt(delta)

    # Evaluate both branches of s
    tsq1 = tau + 2.0 * s1
    tsq2 = tau - 2.0 * s1     # = tau + 2*(-s1)

    # Pick branch maximising |t²| to avoid near-zero division
    if abs(tsq2) > abs(tsq1):
        s = -s1
        tsq = tsq2
    else:
        s = s1
        tsq = tsq1

    if abs(tsq) < 1e-28:
        return None  # degenerate — caller falls through to scipy

    t = cmath.sqrt(tsq)
    inv_t = 1.0 / t

    x00 = (a + s) * inv_t
    x01 = b * inv_t
    x10 = c * inv_t
    x11 = (d + s) * inv_t

    # Verify X·X ≈ A (same tolerance as is_solution)
    r00 = x00 * x00 + x01 * x10
    r01 = x00 * x01 + x01 * x11
    r10 = x10 * x00 + x11 * x10
    r11 = x10 * x01 + x11 * x11

    scale = max(abs(a), abs(b), abs(c), abs(d), 1.0)
    tol = 1e-8 + 1e-5 * scale

    if (abs(r00 - a) > tol or abs(r01 - b) > tol or
            abs(r10 - c) > tol or abs(r11 - d) > tol):
        return None

    return [[x00, x01], [x10, x11]]
--- End Solver ---
Running performance test...
============================= test session starts ==============================
platform linux -- Python 3.12.13, pytest-9.0.2, pluggy-1.6.0
rootdir: /tests
plugins: anyio-4.13.0, jaxtyping-0.3.9
collected 3 items

../tests/test_outputs.py .
--- Performance Summary ---
Validity: True
Total Baseline Time: 27.8706s
Total Solver Time:   11.5810s
Raw Speedup:         2.4066 x
Final Reward (Score): 2.4066
---------------------------
..

=============================== warnings summary ===============================
test_outputs.py: 1100 warnings
  /tests/evaluator.py:79: DeprecationWarning: The `disp` argument is deprecated and will be removed in SciPy 1.18.0.
    X, _ = scipy.linalg.sqrtm(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
==================================== PASSES ====================================
=========================== short test summary info ============================
PASSED ../tests/test_outputs.py::test_solver_exists
PASSED ../tests/test_outputs.py::test_solver_validity
PASSED ../tests/test_outputs.py::test_solver_speedup
================= 3 passed, 1100 warnings in 586.91s (0:09:46) =================
