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

# Discover BLAS thread controllers once at import time. On this host both the
# numpy and scipy OpenBLAS builds default to using all cores (64), which makes
# LAPACK factorizations (eig/schur/sqrtm) extremely slow due to thread
# oversubscription. We limit them to a single thread inside solve().
try:
    from threadpoolctl import ThreadpoolController
    _CONTROLLER = ThreadpoolController()
    _BLAS = [c for c in _CONTROLLER.lib_controllers
             if getattr(c, "user_api", None) == "blas"]
except Exception:
    _BLAS = []

_ORIG = []
for _c in _BLAS:
    try:
        _ORIG.append(_c.num_threads)
    except Exception:
        _ORIG.append(None)


class Solver:
    def __init__(self):
        self._blas = _BLAS
        self._orig = _ORIG

    def solve(self, problem, **kwargs) -> Any:
        A = np.asarray(problem["matrix"], dtype=np.complex128)
        if A.ndim != 2:
            A = np.atleast_2d(A)
        n = A.shape[0]
        if n == 0:
            return {"sqrtm": {"X": []}}

        blas = self._blas
        for c in blas:
            try:
                c.set_num_threads(1)
            except Exception:
                pass
        try:
            X = self._eig_sqrt(A)
            if X is None:
                X = self._scipy_sqrt(A)
            if X is None:
                return {"sqrtm": {"X": []}}
            return {"sqrtm": {"X": X.tolist()}}
        finally:
            for c, o in zip(blas, self._orig):
                if o is not None:
                    try:
                        c.set_num_threads(o)
                    except Exception:
                        pass

    @staticmethod
    def _eig_sqrt(A):
        try:
            w, V = np.linalg.eig(A)
            X = np.linalg.solve(V.T, (V * np.sqrt(w)).T).T
        except Exception:
            return None
        if not np.isfinite(X).all():
            return None
        if np.allclose(X @ X, A, rtol=1e-6, atol=1e-8):
            return X
        return None

    @staticmethod
    def _scipy_sqrt(A):
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                res = scipy.linalg.sqrtm(A, disp=False)
            X = res[0] if isinstance(res, tuple) else res
        except Exception:
            return None
        if X is None or not np.isfinite(X).all():
            return None
        return X
--- 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.8343s
Total Solver Time:   3.9891s
Raw Speedup:         6.9775 x
Final Reward (Score): 6.9775
---------------------------
..

=============================== 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 496.57s (0:08:16) =================
