import torch
import math


def complexity(dstrfs, batch_size=8):
    """
    Measure general complexity of dSTRF function.

    Arguments:
        dstrfs: tensor of dSTRFs with shape [time * channel * lag * frequency]
    
    Returns:
        complexity: nonlinear function complexity, tensor of shape [channel]
    """
    tdim, cdim, ldim, fdim = dstrfs.shape
    if cdim > batch_size:
        return torch.cat([
            complexity(
                dstrfs[:, k*batch_size:(k+1)*batch_size],
                batch_size=batch_size
            ) for k in range(math.ceil(cdim / batch_size))
        ])
    
    dstrfs = dstrfs.transpose(0, 1).reshape(cdim, tdim, ldim*fdim)
    
    sing_vals = torch.linalg.svdvals(dstrfs.float())
    complexty = (sing_vals / sing_vals.max(dim=1, keepdims=True)[0]).sum(dim=1)
    
    return complexty.cpu()