import math
import json
import torch
import librosa
import torchaudio
import os
import numpy as np
import pandas as pd
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
import time


def move_data_to_device(data, device):
    ret = []
    for i in data:
        if isinstance(i, torch.Tensor):
            ret.append(i.to(device))
    return ret

def read_content(filepath):
    '''
    Read the content file for characters, pinyin and tones.

    return:
    dict: {index: [characters, pinyin, tones]}
    exp. {'SS00050001': ['你 好 ', 'ni3 hao3 ', '3 3 ']}
    '''
    res = {}
    with open(filepath, 'r') as f:
        lines = f.readlines()
        for l in lines:
            l = l.replace('\n', ' ').replace('\t', ' ')
            tmp = l.split(' ')
            if len(tmp) == 0:
                break
            number = tmp[0][0:len(tmp[0])-4]
            s = ''
            pinyin = ''
            tones = ''
            for i in range(1, len(tmp)):
                if len(tmp[i]) == 0:
                    continue
                # need blank space or not?
                if i % 2 == 0:
                    pinyin += tmp[i] + ' '
                    tones += tmp[i][-1] + ' '
                else:
                    s += tmp[i] + ' '
            res[number] = [s, pinyin, tones]
    return res

def read_dataset_index(filepath='/kaggle/input/paddle-speech/AISHELL-3/train'):
    '''
    get all audio files' index and file paths
    read content.txt to get corresponding words, pinyin, tones, duration

    return dataframe:
    ['index', 'filepath', 'word', 'pinyin', 'tone', 'duration']

    5 tones in total, 5 represents neutral tone
    '''
    features = read_content(os.path.join(filepath, 'content.txt'))

    start_time = time.time()
    count = 0

    durations = {}
    with open('/kaggle/input/durations/durations.txt', 'r') as f:
        lines = f.readlines()
        for l in lines:
            tmp = (l.replace('\n', '')).split(' ')
            if len(tmp) != 0:
                durations[tmp[0]] = float(tmp[1])

    audio_path = os.path.join(filepath, 'wav')#这里要删掉
    indexes = []
    for root, dirs, files in os.walk(audio_path):
        for f in files:
            if f.endswith('.wav'):
                count += 1
                index = f[0:len(f)-4]
                filepath = os.path.join(audio_path, index[0:len(index)-4], f)
                word, py, tone = features[index]
                # du = librosa.get_duration(filename=filepath)
                du = durations[index]
                indexes.append((index, filepath, word, py, tone, du))
    
    end_time = time.time()
    print('#wav file read:', count)
    print('read dataset index time: ', end_time - start_time)

    '''indexes = sorted(indexes, key=lambda x: x[0])
    with open('./durations.txt', 'w') as f:
        for i in indexes:
            f.write(i[0]+ ' ' + str(i[5]) + '\n')'''

    return pd.DataFrame.from_records(indexes, columns=['index', 'filepath', 'word', 'pinyin', 'tone', 'duration'])


def read_dataset_index(filepath='/kaggle/input/paddle-speech/AISHELL-3/train'):
    '''
    get all audio files' index and file paths
    read content.txt to get corresponding words, pinyin, tones, duration

    return dataframe:
    ['index', 'filepath', 'word', 'pinyin', 'tone', 'duration']

    5 tones in total, 5 represents neutral tone
    '''
    features = read_content(os.path.join(filepath, 'content.txt'))

    start_time = time.time()
    count = 0

    durations = {}
    with open('/kaggle/input/durations/durations.txt', 'r') as f:
        lines = f.readlines()
        for l in lines:
            tmp = (l.replace('\n', '')).split(' ')
            if len(tmp) != 0:
                durations[tmp[0]] = float(tmp[1])

    audio_path = os.path.join(filepath, 'wav')#这里要删掉
    indexes = []
    for root, dirs, files in os.walk(audio_path):
        for f in files:
            if f.endswith('.wav'):
                count += 1
                index = f[0:len(f)-4]
                filepath = os.path.join(audio_path, index[0:len(index)-4], f)
                word, py, tone = features[index]
                # du = librosa.get_duration(filename=filepath)
                du = durations[index]
                indexes.append((index, filepath, word, py, tone, du))
    
    end_time = time.time()
    print('#wav file read:', count)
    print('read dataset index time: ', end_time - start_time)

    '''indexes = sorted(indexes, key=lambda x: x[0])
    with open('./durations.txt', 'w') as f:
        for i in indexes:
            f.write(i[0]+ ' ' + str(i[5]) + '\n')'''

    return pd.DataFrame.from_records(indexes, columns=['index', 'filepath', 'word', 'pinyin', 'tone', 'duration'])


def collate_fn(batch):
    inp = []
    f0 = []
    word = []
    tone = []
    max_frame_num = 1600
    for sample in batch:
        max_frame_num = max(max_frame_num, sample[0].shape[0], sample[1].shape[0], sample[2].shape[0], sample[3].shape[0])
    for sample in batch:
        inp.append(
            torch.nn.functional.pad(sample[0], (0, 0, 0, max_frame_num - sample[0].shape[0]), mode='constant', value=0))
        f0.append(
            torch.nn.functional.pad(sample[1], (0, max_frame_num - sample[1].shape[0]), mode='constant', value=0))
        word.append(
            torch.nn.functional.pad(sample[2], (0, 50 - sample[2].shape[0]), mode='constant', value=0))
        tone.append(
            torch.nn.functional.pad(sample[3], (0, 50 - sample[3].shape[0]), mode='constant', value=0))
    inp = torch.stack(inp)
    f0 = torch.stack(f0)
    word = torch.stack(word)
    tone = torch.stack(tone)

    return inp, f0, word, tone


def get_data_loader(split, args):
    Dataset = MyDataset(
        dataset_root=args['dataset_root'],
        split=split,
        sampling_rate=args['sampling_rate'],
        sample_length=args['sample_length'],
        frame_size=args['frame_size'],
    )
    Dataset.dataset_index=Dataset.dataset_index[:32]
    Dataset.index=Dataset.index[:32]
    data_loader = DataLoader(
        Dataset,
        batch_size=args['batch_size'],
        num_workers=args['num_workers'],
        pin_memory=True,
        shuffle=True, # changed into True cuz audio files recorded by same speaker are stored in the same folder
        collate_fn=collate_fn,
    )
    return data_loader

class MyDataset(Dataset):
    def __init__(self, dataset_root, split, sampling_rate, sample_length, frame_size):
        self.dataset_root = dataset_root
        self.split = split # train or test
        self.sampling_rate = sampling_rate
        self.sample_length = sample_length
        self.frame_size = frame_size
        self.frame_per_sec = int(1 / self.frame_size)

        # self.annotations = get_annotations(get_all_file_names(os.path.join(self.dataset_root, 'AISHELL-3', split)), level='word')

        self.dataset_index = read_dataset_index(os.path.join(self.dataset_root, 'AISHELL-3', split)) # maybe can be removed

        self.duration = {}
        self.index = self.index_data()
        # print(len(self.index))
        self.dataset_index=self.dataset_index[:10]
        self.index=self.index[:10]
        
        self.pinyin = {} # read encoded pinyin
        with open('/kaggle/input/pinyin-encode/pinyin.txt', 'r') as f:
            lines = f.readlines()
            i = 0
            for l in lines:
                self.pinyin[l.replace('\n', '')] = i
                i += 1

    
    def index_data(self):
        '''
        Prepare the index for the dataset, i.e., the audio file name and starting time of each sample

        go through self.dataset_index to get duration and then calculate
        '''
        # duration already in dataset_index
        # TODO
        # pass
        index = []
        for indexs, row in self.dataset_index.iterrows():
            duration = row['duration']
            num_seg = math.ceil(duration / self.sample_length)
            for i in range(num_seg):
                # index.append([row['index'], i * self.sample_length])
                index.append([indexs, i * self.sample_length])
            self.duration[row['index']] = row['duration']

        return index

    
    def __len__(self):
        return len(self.index)
    
    def __getitem__(self, idx):
        '''
        int idx: index of the audio file (not exp.SSB00050001)

        return mel spectrogram, FUNDAMENTAL FREQUENCY(crepe/pyin), words, tones
        '''
        audio_fn, start_sec = self.index[idx]
        end_sec = start_sec + self.sample_length
        # print(start_sec, end_sec)
        #???
        audio_fp = self.dataset_index.loc[audio_fn,'filepath']
        # audio_fp = jpath('./dataset/AISHELL-3/train/wav/SSB0005/SSB0005',audio_fp,'.wav')
        #/kaggle/input/paddle-speech/AISHELL-3/train/wav/SSB0005/SSB00050001.wav

        # TODO: calculate mel spectrogram
        mel = None
        #load data from file
        waveform, sample_rate = torchaudio.load(audio_fp)
        waveform = torchaudio.transforms.Resample(sample_rate, self.sampling_rate)(waveform)
        mel_spec = torchaudio.transforms.MelSpectrogram(sample_rate=self.sampling_rate, n_fft=2048, hop_length=100, n_mels=256)(waveform)
        mel_spec = torch.mean(mel_spec,0)
        # print(mel_spec.shape)    

        # TODO: calculate fundamental frequency
        f0 = None
        waveform, sr = librosa.load(audio_fp, sr=self.sampling_rate)
        f0 = torch.from_numpy(librosa.yin(waveform, fmin=50, fmax=550, hop_length=100))

        # get labels???
        # word_roll, tone_roll = self.get_labels(self.annotations[self.dataset_index.loc[audio_fn, 'index']], self.dataset_index.loc[audio_fn,'duration'])
        words = self.dataset_index.loc[audio_fn, 'pinyin']
        w = words.split(' ')
        word_roll = []
        for i in range(0, len(w)):
            
            if len(w[i]) != 0:
                if self.pinyin.get(w[i][0:-1]) == None:
                    self.pinyin[w[i][0:-1]] = len(self.pinyin)
                word_roll.append(self.pinyin[w[i][0:-1]])
        tones = self.dataset_index.loc[audio_fn, 'tone']
        t = tones.split(' ')
        tone_roll = []
        for tone in t:
            if len(tone) != 0:
                tone_roll.append(int(tone))
        
        spectrogram_clip = None
        f0_clip = None
        onset_clip = None
        offset_clip = None
        word_clip = None
        tone_clip = None

        # TODO: create clips
        start_frame = int(start_sec * self.frame_per_sec)
        end_frame = start_frame + 1600 #int(end_sec * self.frame_per_sec)
        # print(start_frame, end_frame)
        spectrogram_clip = mel_spec[:, start_frame:end_frame].T
        f0_clip = f0[start_sec:end_sec]
        #word_clip = word_roll[start_frame:end_frame]
        #tone_clip = tone_roll[start_frame:end_frame]

        # print(tone_roll)
        #return spectrogram_clip, f0_clip, onset_clip, offset_clip, pinyin_clip, tone_clip
        return spectrogram_clip, f0_clip, torch.Tensor(word_roll), torch.Tensor(tone_roll) #word_clip, tone_clip
    
    def get_labels(self, annotation_data, duration):
        '''
        This function read annotation from file, and then convert annotation from note-level to frame-level
        Because we will be using frame-level labels in training.
        '''
        # TODO
        # pass
        frame_num = math.ceil(duration * self.frame_per_sec)

        word_roll = torch.zeros(size=(frame_num + 1,), dtype=torch.long)
        tone_roll = torch.zeros(size=(frame_num + 1,), dtype=torch.long)
        # f0_roll = torch.zeros(size=(frame_num + 1,), dtype=torch.long)
        # mel_roll = torch.zeros(size=(frame_num + 1,), dtype=torch.long)
        for note in annotation_data:
            start_time, end_time, mark = note  # Assuming annotation format: (start_time, end_time, pitch)

            # Convert note start and end times to frame indices
            start_frame = int(start_time * self.frame_per_sec)
            end_frame = int(end_time * self.frame_per_sec)

            # Clip frame indices to be within the valid range, no need in this task
            start_frame = max(0, min(frame_num, start_frame))
            end_frame = max(0, min(frame_num, end_frame))
            #print(start_frame, end_frame)

            # WORD LEVEL Mark the frames corresponding to the note
            word_roll[start_frame:end_frame+1] = self.pinyin[mark[:-1]] #mark[:-1]
            tone_roll[start_frame:end_frame+1] = int(mark[-1])
        # print(tone_roll)
        return word_roll, tone_roll