import torch
import math


def gain_change(dstrfs, batch_size=8):
    """
    Measure standard deviation of dSTRF gains.
    
    Arguments:
        dstrfs: tensor of dSTRFs with shape [time * channel * lag * frequency]
    
    Returns:
        gain_change: shape change parameter, tensor of shape [channel]
    """
    tdim, cdim, ldim, fdim = dstrfs.shape
    if cdim > batch_size:
        return torch.cat([
            gain_change(
                dstrfs[:, k*batch_size:(k+1)*batch_size],
                batch_size=batch_size
            ) for k in range(math.ceil(cdim / batch_size))
        ])
    
    return dstrfs.norm(dim=[-2, -1]).std(dim=0).cpu()