import os
import sys
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import qgrid
import xml.etree.ElementTree as ET
from torch.utils.data import Dataset, DataLoader, Subset
from torch.nn.utils.rnn import pad_sequence
import librosa
from torch.nn.utils.rnn import pack_sequence
from torch.utils.data import DataLoader
import torch
from sklearn.model_selection import train_test_split
import madmom
import mir_eval as me
import c_evaluation

data = pd.read_csv('all_data_list.csv')

anot_dataset_path = "Annotations/" #Path to original annotations
wav_dataset_path = "Wavs/" #Path to original wavs
atar_dataset_path = "processed_annotations/" #Path to preprocessed annotations
tar_dataset_path = "processed_wavs/" #Path to preprocessed wavs


def get_filename(index):
	file = data.at[index, 'filename']
	return file


def get_audio(filename):  # get audio signal
	b = wav_dataset_path + filename + '.wav'
	signal, sr = librosa.load(b, sr=44100)
	return signal, sr


def get_input(index):  # get input for NN
	filename = get_filename(index)
	with open(tar_dataset_path + filename + '.npy', 'rb') as f:
		input = np.load(f)
		input[60:199, :] = input[60:199, :] * 2  # multiplying median spectrogram part showed better results
	input = np.array(input)
	input = torch.tensor(input)
	input = torch.transpose(input, -1, 0)
	return input


def get_inverse_label(label):  # get original beat times
	label = np.where(label == 1)
	label = np.array(label) / 100
	# label[label == 0.5] = 0
	return label


def get_label_tensor(index):  # get label for NN
	filename = get_filename(index)
	with open(atar_dataset_path + filename + '.npy', 'rb') as f:
		label = np.load(f)
	label = np.array(label)
	label = torch.tensor(label)
	label = torch.transpose(label, -1, 0)
	return label


def get_label(index):  # get label
	filename = get_filename(index)
	with open(atar_dataset_path + filename + '.npy', 'rb') as f:
		label = np.load(f)
	label = np.array(label)
	return label


def sec_to_frame(num_frames, anotace):   # transform beat times for NN
	frame_pos = np.zeros(num_frames)
	for idx, val in enumerate(frame_pos):
		if idx in anotace:
			frame_pos[idx] = 1
			if idx + 1 < num_frames:
				frame_pos[idx + 1] = 0.5
				frame_pos[idx - 1] = 0.5
	labels = np.array(frame_pos)
	return labels


def anot_to_vec(anotace, frame_rate=100):  #transform txt to vector
	label = list()
	with open(anotace) as anot:
		for line in anot:
			label.append(line.strip())
	label = np.array(label, dtype=np.float32)
	label = label * frame_rate
	label = np.ceil(label)
	return label


def show_spec(input):  #show spectogram
	plt.imshow(input, origin='lower', aspect='auto', cmap='jet')
	plt.colorbar(format='%.1f')
	plt.xlabel('Časový segment [-]')
	plt.ylabel('Řadové číslo spektrogramu [-]')
	plt.yticks(ticks=[0, 20, 40, 60, 80, 100, 120], labels=[0, 1, 2, 3, 4, 5, 6])
	plt.tight_layout()


def show_spec_label(input, label):  #show spectrogram with labels
	plt.imshow(input, origin='lower', aspect='auto', cmap='jet')
	plt.colorbar(format='%.1f')
	plt.xlabel('Časový segment [-]')
	plt.ylabel('Řadové číslo spektrogramu [-]')
	plt.yticks(ticks=[0, 20, 40, 60, 80, 100, 120], labels=[0, 1, 2, 3, 4, 5, 6])
	plt.tight_layout()
	for idx, val in enumerate(label):
		if val == 1:
			plt.axvline(x=idx, color='w', linestyle=':', linewidth=3)


class Data(Dataset):  #class for dataloader
	def __init__(self, data):
		self.data = data
		self.indices = data.index

	def __getitem__(self, i):
		return get_input(self.indices[i]), get_label_tensor(self.indices[i])

	def __len__(self):
		return len(self.data)


def collate_fn(batch):  # pad sequence for batch sizes > 1
	features, labels = zip(*batch)
	features = pad_sequence(features)
	labels = pad_sequence(labels)
	return features, labels


def output_feature(model, feature):  # output from given preprocessed data
	feature = feature.permute(1, 2, 0)
	feature = feature.cuda()
	feature = feature.float()
	with torch.no_grad():
		out = model(feature)
	out = out.cpu()
	a = np.array(out)
	a = a[0, 1, :]
	a = np.exp(a)
	return a


def output_index(model, index):  # output from given index
	feature = get_input(index)
	feature = feature.unsqueeze(0)
	feature = feature.permute(0, 2, 1)
	feature = feature.cuda()
	feature = feature.float()
	with torch.no_grad():
		out = model(feature)
	out = out.cpu()
	a = np.array(out)
	a = a[0, 1, :]
	a = np.exp(a)
	return a


def cross_val_indices(k, cvi_filename, data_path):  # generate cross val folds
	data_fun = pd.read_csv(data_path)
	a = data_fun['dataset']

	be_indices = np.array(data_fun.index[(a == 'BE') == True].tolist())
	br_indices = np.array(data_fun.index[(a == 'BR') == True].tolist())
	gt_indices = np.array(data_fun.index[(a == 'GT') == True].tolist())
	hw_indices = np.array(data_fun.index[(a == 'HW') == True].tolist())
	sm_indices = np.array(data_fun.index[(a == 'SM') == True].tolist())

	np.random.shuffle(be_indices)
	np.random.shuffle(br_indices)
	np.random.shuffle(gt_indices)
	np.random.shuffle(hw_indices)
	np.random.shuffle(sm_indices)

	be = np.array_split(be_indices, k)
	br = np.array_split(br_indices, k)
	gt = np.array_split(gt_indices, k)
	hw = np.array_split(hw_indices, k)
	sm = np.array_split(sm_indices, k)

	fold = []
	for i in range(k):
		fold.append(np.concatenate((be[i], br[i], gt[i], hw[i], sm[i])))
	flat_fold = np.sort([item for sublist in fold for item in sublist]).tolist()
	set1 = set(flat_fold)
	train_indices = []
	val_indices = []
	tt_ratio = 1/(k-1)

	for i in range(k):
		set2 = set(fold[i])
		subtrainset = list(set1 - set2)
		subtrainset, subvalset = train_test_split(subtrainset, shuffle=True, test_size=tt_ratio)
		train_indices.append(subtrainset)
		val_indices.append(subvalset)

	test_indices = fold
	cross_val_folds = pd.DataFrame(fold)
	cross_val_folds.to_csv(cvi_filename + ".csv")

	return train_indices, val_indices, test_indices


def load_cross_val_indices(file_path):  # load cross val folds

	data = pd.read_csv(file_path)
	fold = np.array(data)
	flat_fold = np.sort([item for sublist in fold for item in sublist])
	flat_fold = np.unique(flat_fold)
	flat_fold = flat_fold[np.logical_not(np.isnan(flat_fold))]

	train_indices = []
	val_indices = []
	test_indices = []
	tt_ratio = 1 / (fold.shape[0] - 1)
	set1 = set(flat_fold)

	for i in range(fold.shape[0]):
		set2 = set(fold[i])
		subtrainset = list(set1 - set2)
		subtrainset, subvalset = train_test_split(subtrainset, shuffle=True, test_size=tt_ratio)
		train_indices.append(subtrainset)
		val_indices.append(subvalset)
		test_fold = fold[i][1:]
		test_indices.append(np.unique(test_fold[np.logical_not(np.isnan(test_fold))]))

	return train_indices, val_indices, test_indices


def process_batch_tcn(model, feature, label, loss_function, device):  # process batch with TCN
	feature = feature.float()
	feature = feature.permute(1, 2, 0)
	label = label.transpose(1, 0)
	out = model(feature.to(device))
	label = torch.tensor(label, dtype=torch.long, device=device)
	loss = loss_function(out, label)
	return loss


def process_batch_cnntcn(model, feature, label, loss_function, device):  # process batch with CNNTCN
	feature = feature.float()
	feature = feature.permute(1, 0, 2)
	label = label.transpose(1, 0)
	out = model(feature.to(device))
	label = torch.tensor(label, dtype=torch.long, device=device)
	loss = loss_function(out, label)
	return loss


def process_batch_rnn(model, feature, label, loss_function, device):  # process batch with BLSTM or BGRU
	feature = feature.float()
	out = model(feature)
	out = out.permute(1, 2, 0)
	label = label.transpose(1, 0)
	label = torch.tensor(label, dtype=torch.long, device=device)
	loss = loss_function(out, label)
	return loss

