import numpy as np
import torch
import torch.nn as nn


class Subspectrogram(object):
    def __init__(self, height, hop_size):
        self.height = height
        self.hop_size = hop_size

    def __call__(self, sample):
        if len(sample.shape) < 3:
            sample = sample.reshape(1, *sample.shape)
        # sample shape: 1 x num_mels x num_frames
        sub_specs = []
        neg = 0 if self.height == self.hop_size else self.hop_size
        for i in range(0, sample.shape[1]-neg, self.hop_size):
            sub_spec = sample[:, i:i+self.height:,]
            sub_specs.append(sub_spec)
        return np.concatenate(sub_specs)



if __name__ == '__main__':
    import numpy as np
    sub_spec_tnfm = Subspectrogram(20, 20)
    X = np.random.rand(1, 80, 40)
    Y = sub_spec_tnfm(X)
    print(f'\tSub-Spectrogram transformation from shape {X.shape} to {Y.shape}')
    print('\tDone ...')