import torch

def mean_centroid(x, sr=1.0, patch_size=16):
    # x: nvar, L
    f = torch.fft.rfft(x, dim=-1).abs()
    freqs = torch.fft.rfftfreq(x.size(-1), 1/sr).to(x.device)
    return int(((1 / (((f * freqs).sum(-1) / f.sum(-1)).mean())) // patch_size) * patch_size)

def generate_reservoir_matrix(n, sparsity=0.05, spectral_radius=0.9, seed=None):
    if seed is not None:
        torch.manual_seed(seed)
    
    # Step 1: Random matrix with values in [-1, 1]
    W = torch.rand(n, n) * 2 - 1

    # Step 2: Apply sparsity mask
    mask = (torch.rand(n, n) < sparsity).float()
    W *= mask

    # Step 3: Normalize to desired spectral radius
    eigenvalues = torch.linalg.eigvals(W).abs()
    max_eigenvalue = torch.max(eigenvalues)
    if max_eigenvalue > 0:
        W *= spectral_radius / max_eigenvalue

    return W