--- Final Solver (used for performance test) ---
import numpy as np
from scipy.linalg import schur
from scipy.linalg._matfuncs_sqrtm import _sqrtm_triu
from scipy.linalg import sqrtm as scipy_sqrtm
from typing import Any
import numba


@numba.njit(cache=True)
def _denman_beavers(A, n, max_iter=50):
    """Denman-Beavers迭代计算矩阵平方根。"""
    Y = A.copy()
    Z = np.eye(n, dtype=numba.complex128)
    for k in range(max_iter):
        Y_inv = np.linalg.inv(Y)
        Z_inv = np.linalg.inv(Z)
        Y_new = 0.5 * (Y + Z_inv)
        Z_new = 0.5 * (Z + Y_inv)
        diff = 0.0
        for ii in range(n):
            for jj in range(n):
                d = abs(Y_new[ii, jj] - Y[ii, jj])
                if d > diff:
                    diff = d
        Y = Y_new
        Z = Z_new
        if diff < 1e-12:
            break
    return Y


class Solver:
    def __init__(self):
        # JIT预热
        _n = 3
        _A = np.eye(_n, dtype=np.complex128)
        _denman_beavers(_A, _n, 5)

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

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

        # 小矩阵使用DB迭代
        if n <= 20:
            X = self._sqrtm_db(A, n)
            if X is not None:
                return {"sqrtm": {"X": X.tolist()}}

        # 大矩阵使用Schur + _sqrtm_triu + 反变换
        X = self._sqrtm_direct(A, n)
        if X is not None:
            return {"sqrtm": {"X": X.tolist()}}

        # 回退到scipy
        X = self._sqrtm_scipy(A)
        if X is not None:
            return {"sqrtm": {"X": X.tolist()}}

        return {"sqrtm": {"X": []}}

    def _sqrtm_db(self, A, n):
        """Denman-Beavers迭代。"""
        try:
            X = _denman_beavers(A.astype(np.complex128), n, 50)
            if not np.all(np.isfinite(X)):
                return None
            return X
        except Exception:
            return None

    def _sqrtm_direct(self, A, n):
        """直接Schur + _sqrtm_triu + 反变换。"""
        try:
            Ac = A.astype(np.complex128)
            T, Z = schur(Ac, output='complex')
            R = _sqrtm_triu(T)
            X = Z @ R @ Z.conj().T
            return X
        except Exception:
            return None

    def _sqrtm_scipy(self, A):
        """回退到scipy.linalg.sqrtm。"""
        try:
            X, _ = scipy_sqrtm(A, disp=False)
            return X
        except Exception:
            return None
--- 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: 35.7835s
Total Solver Time:   11.5451s
Raw Speedup:         3.0995 x
Final Reward (Score): 3.0995
---------------------------
..

=============================== 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 686.43s (0:11:26) =================
