import pandas as pd

import numpy as np

import librosa

from sklearn import neighbors

from multiprocessing import Pool

import warnings

import os

import tqdm.notebook as tqdm

import glob

import unidecode

def feature_extractor(file):

    warnings.filterwarnings('ignore')

    x, fs = librosa.load(file)

    feat =  librosa.feature.mfcc(x)

    warnings.filterwarnings("default")

    return np.concatenate((feat.mean(axis=1), feat.min(axis=1), feat.max(axis=1), feat.std(axis=1)))

def normalize_str(txt):

    invalids = set()

    # vocabulary

    valid_chars = (" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",

                   "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'") # New valid characters 

    # lowercase

    new_txt = txt.lower().strip()

    # remove characters not in vocabulary

    res_arr = []

    for c in new_txt:

        if c in valid_chars:

            res_arr.append(c)

        else:

            #remove accent and see if it is valid

            non_accent_c = unidecode.unidecode(c)

            if non_accent_c in valid_chars:

                res_arr.append(non_accent_c)

            else:

                # a character we don't know

                invalids.update(c)

                res_arr.append(' ')

    res = ''.join(res_arr).strip()

    return ' '.join(res.split())

datasetdir = "/data/CV_unpacked/cv-corpus-7.0-2021-07-21/lg"

df = pd.read_csv(f'{datasetdir}/validated.tsv', delimiter="\t")

X = [feature_extractor(f"{datasetdir}/clips/{_.path}") for i,_ in tqdm.tqdm(df.iterrows())]

#That took a while, so save for later.

np.save("feats", X)

X = np.load("feats.npy")

# Assign a unique number to each utterance and create lookup.

y = df.index

labels = [normalize_str(_) for _ in df.sentence]

clf = neighbors.KNeighborsClassifier(n_neighbors=1, leaf_size=100)

clf.fit(X, y)

# Generate test predictions

testdir = "/data/test_audio/test_audio"

testfiles = glob.glob(f"{testdir}/*.mp3")

testfeats = [feature_extractor(_) for _ in tqdm.tqdm(testfiles)]

testpreds = clf.predict(testfeats)

testtrans = [labels[_].replace("'","") for _ in testpreds]

with open("predictions.csv", "w") as f:

    f.write("Clip_ID,Target\n")

    for file, transcription in zip(testfiles,testtrans):

        ID = os.path.split(os.path.splitext(file)[0])[1]

        f.write(f"{ID},{transcription}\n")