--- Final Solver (used for performance test) ---
"""
Fast matrix square root solver.

Avoids the input validation overhead of scipy.linalg.sqrtm by directly
calling LAPACK routines for the Schur decomposition and Sylvester solves.
"""
import numpy as np
import scipy.linalg
import scipy.linalg.lapack as lapack
from scipy.linalg.lapack import ztrsyl, dtrsyl


class Solver:
    def __init__(self):
        # Pre-bind LAPACK functions for speed
        self._zgees = lapack.zgees
        self._ztrsyl = ztrsyl
        self._dtrsyl = dtrsyl
        # Use scipy's optimized sqrtm of triangular matrix
        from scipy.linalg._matfuncs_sqrtm import _sqrtm_triu
        self._sqrtm_triu = _sqrtm_triu

    def solve(self, problem, **kwargs):
        A = problem["matrix"]
        if not isinstance(A, np.ndarray):
            A = np.asarray(A)

        n = A.shape[0] if A.ndim == 2 else 0
        if n == 0:
            return {"sqrtm": {"X": []}}
        if A.shape != (n, n):
            return {"sqrtm": {"X": []}}

        if n == 1:
            a = A[0, 0]
            r = np.sqrt(a)
            return {"sqrtm": {"X": [[r]]}}

        is_complex = np.iscomplexobj(A)

        try:
            if is_complex:
                if A.dtype != np.complex128:
                    A = A.astype(np.complex128)
                # Direct LAPACK Schur decomposition
                Af = np.asfortranarray(A)
                T, _, _, Z, _, info = self._zgees(
                    lambda x: 1, Af, compute_v=1, sort_t=0, overwrite_a=1
                )
                if info != 0:
                    return {"sqrtm": {"X": []}}
                R = self._sqrtm_triu(T)
                X = Z @ R @ Z.conj().T
            else:
                if A.dtype != np.float64:
                    A = A.astype(np.float64)
                # Use real Schur for real matrices
                from scipy.linalg.lapack import dgees
                Af = np.asfortranarray(A)
                T, _, _, Z, _, info = dgees(
                    lambda x: 1, Af, compute_v=1, sort_t=0, overwrite_a=1
                )
                if info != 0:
                    return {"sqrtm": {"X": []}}
                R = self._sqrtm_triu(T)
                X = Z @ R @ Z.T
        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: 34.2724s
Total Solver Time:   12.0843s
Raw Speedup:         2.8361 x
Final Reward (Score): 2.8361
---------------------------
..

=============================== 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 663.47s (0:11:03) =================
