import torch


def mc_stft(y_s, n_fft, hop_length, win_length):
    """
    Multi-Channel STFT

    Shape:
        y_s: [B, C, T]

    Returns:
        complex_value: [B, C, F, T]
    """
    assert y_s.dim() == 3
    batch_size, num_channels, num_wav_samples = y_s.size()

    # [B * C, F, T] in C
    stft_coefficients = torch.stft(
        y_s.reshape(batch_size * num_channels, num_wav_samples),  # [B * C, T]
        n_fft=n_fft,
        hop_length=hop_length,
        window=torch.hann_window(win_length, device=y_s.device),
        win_length=win_length,
        return_complex=True
    )

    return stft_coefficients.reshape(batch_size, num_channels, stft_coefficients.shape[-2], stft_coefficients.shape[-1])