import torch


def stft_wrapper(x, fft_n, frame_shift, frame_length, window, 
                 pad_mode="constant"):
    """Due to the different signature of torch.stft, write a 
    wrapper to handle this

    input
    -----
      x: tensor, waveform, (batch, length)
      window: tensor, window coef, (frame_length, )
    
    output
    ------
      tensor: (batch, frame_num, bin_num, 2)
    
    """
    # there are better ways, but for convenience
    if torch.__version__.split('.')[1].isnumeric() and \
       int(torch.__version__.split('.')[1]) < 7:
        #  torch 1.6.*
        return torch.stft(x, fft_n, frame_shift, frame_length, 
                          window=window, onesided=True, pad_mode=pad_mode)
    else:
        #  torch > 1.7
        return torch.stft(x, fft_n, frame_shift, frame_length, 
                          window=window, onesided=True, pad_mode=pad_mode,
                          return_complex=False)