from common.dataset import AudioDataset, get_data_loader
from common.features import  FilterbankFeatures
from common import audio, helpers


import torch
import numpy as np
from jasper import config

cfg = config.load('configs/jasper10x5dr_speedp-online_speca.yaml')
train_dataset_kw, train_features_kw = config.input(cfg, 'train')

val_dataset_kw, val_features_kw = config.input(cfg, 'val')
# dataset_dir = 'test_small/test_small'
dataset_dir = '/data2/hdu/zlz/dataset'
val_manifests = ['test_small/test_small/dev-clean-wav.json']
train_manifest = ['/data2/hdu/zlz/dataset/librispeech-train-clean-100-wav.json',
                  '/data2/hdu/zlz/dataset/librispeech-train-clean-360-wav.json',
                  '/data2/hdu/zlz/dataset/librispeech-train-other-500-wav.json']

symbols = helpers.add_ctc_blank(cfg['labels'])
# print(train_features_kw)
# print(train_dataset_kw)

train_dataset_kw['sort_by_duration'] = True
val_dataset_kw['max_duration'] = 16.7

train_features_kw['spec_augment'] = None
val_features_kw['pad_to_max_duration'] = True
train_dataset_kw['pad_to_max_duration'] = True



# print(train_features_kw)
val_dataset = AudioDataset(dataset_dir,
                           train_manifest,
                           symbols,
                           **train_dataset_kw)


train_feat_proc = FilterbankFeatures(**train_features_kw)
seq_len = torch.zeros(1, dtype=torch.int32)

for i in range(1):
    print(i)
    a, b, _, __ = val_dataset[i]
    # a = np.expand_dims(a, axis=0)
    audio1 = torch.zeros(1, a.shape[0])
    audio1[0].narrow(0, 0, a.shape[0]).copy_(a)
    seq_len[0] = b
    c, d  = train_feat_proc(audio1, seq_len)
    print("c shape:", c.shape)

# print(val_dataset[0])