--- Final Solver (used for performance test) ---
from contextlib import nullcontext
from typing import Any

import numpy as np
import scipy.linalg

try:
    from threadpoolctl import ThreadpoolController

    _CONTROLLER = ThreadpoolController()
except Exception:  # pragma: no cover - threadpoolctl should always be present
    _CONTROLLER = None


class Solver:
    def __init__(self) -> None:
        # The benchmark environment links OpenBLAS with up to 64 threads and
        # NO_AFFINITY on a many-core box. For the matrix sizes seen here, the
        # LAPACK kernels behind eig/schur suffer severe thread oversubscription
        # (a decomposition that costs ~12ms with one thread balloons to seconds
        # with 64 threads). Running with a single BLAS thread avoids it.
        #
        # The limit is scoped to our own computation via a context manager (the
        # controller is cached to keep per-call overhead low) rather than
        # changed globally, so the limit cannot leak into any other timing.
        self._limit = _CONTROLLER

    def solve(self, problem: dict, **kwargs) -> Any:
        A = np.asarray(problem["matrix"], dtype=complex)
        ctx = self._limit.limit(limits=1) if self._limit is not None else nullcontext()

        with ctx:
            # Principal square root via eigendecomposition: A = V diag(w) V^-1,
            # so X = V diag(sqrt(w)) V^-1 satisfies X @ X = A. np.sqrt takes the
            # principal branch (non-negative real part) -> principal root.
            try:
                w, V = np.linalg.eig(A)
                sw = np.sqrt(w)
                # X = V @ diag(sw) @ inv(V), without forming inv(V):
                #   X.T solves V.T @ X.T = (V * sw).T
                X = np.linalg.solve(V.T, (V * sw).T).T

                # eig loses accuracy when V is ill-conditioned (defective or
                # near-defective A). Verify the reconstruction (tighter than the
                # checker's rtol=1e-5) and fall back if it is not good enough.
                if np.all(np.isfinite(X)) and np.allclose(
                    X @ X, A, rtol=1e-6, atol=1e-9
                ):
                    return {"sqrtm": {"X": X.tolist()}}
            except Exception:
                pass

            # Robust fallback: the same Schur-based algorithm the reference
            # uses, still single-threaded so it stays fast. Mirror the
            # reference's behavior exactly so this path is never worse than it.
            try:
                X, _ = scipy.linalg.sqrtm(A, disp=False)
            except Exception:
                return {"sqrtm": {"X": []}}
            return {"sqrtm": {"X": X.tolist()}}
--- 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.3584s
Total Solver Time:   4.1790s
Raw Speedup:         6.5467 x
Final Reward (Score): 6.5467
---------------------------
..

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