import tensorflow as tf
import librosa
import pandas as pd
import os
import unicodedata
import soundfile
import logging
import re
from tqdm import tqdm
from glob import glob
from sklearn.model_selection import train_test_split
from random import choice

logging.basicConfig(level=logging.NOTSET)
FLAGS = tf.compat.v1.app.flags.FLAGS

patterns = [r'(\.)+', r'(\([a-zA-Z0-9 _\(\)]+\))+', r'(<YWN)+', r'(?<=\s)=+\b', r'\b=+', r'\d(?=\w+)', r'(?<=\w)+\d', r'%(?=\w)+',
            r'\[(?=((\s|\'|\w|=|~|@|<|\-|\)|\()+))', r'(?<=(\s|\'|\w|=|~|@|>|\-|\)))+\]', r'(?<=\s)\-+', r'<[\w@%]+', r'[\w@%]+>', r'@+']


def main(_):
    source_dir = FLAGS.source_dir
    transcripts = glob(source_dir + "/transcripts/TRN/*.trn")
    clip_range = range(5, 10)
    data = []
    file_prefix = "sbc-"
    file_index = 1
    with tqdm(total=len(transcripts)) as bar:
        for idx, transcript_file in enumerate(transcripts):
            print(transcript_file)
            df = open(transcript_file, "r", encoding="utf-8")
            clip_duration = choice(clip_range)
            clip_transcript = ''
            curr_start = 0
            curr_start_idx = 0
            df_transcripts = pd.DataFrame(columns=["start", "end", "transcript", "selected clip_duration", "actual clip duration"])
            df_index = 0
            audio_file = source_dir + "/clips/" + transcript_file.split("/")[-1].split(".trn")[0] + ".wav"
            audio_data, sr = librosa.load(audio_file, sr=None)
            audio_duration = librosa.get_duration(audio_data)
            for i, line in enumerate(df):
                line = line.replace("\n", "")
                if line.strip() == "":
                    continue
                curr_transcript = line.split("\t")[-1]
                time_map = re.split(r'(\t|\s)+', line)
                start, end = time_map[0], time_map[2]
                if start == "":
                    start = time_map[2]
                    end = time_map[4]
                try:
                    start = float(start)
                    end = float(end)
                except:
                    print(f'{time_map}, {line}')
                if pd.isnull(curr_transcript):
                    curr_transcript = ''

                curr_transcript = re.sub(r'(' + "|".join(patterns) + r')', '', curr_transcript)
                curr_transcript = re.sub(r'((?<=\s)=+\b|\b=+|,|\?|~)', '', curr_transcript)
                curr_transcript = re.sub(r'\s+', ' ', curr_transcript)
                curr_transcript = curr_transcript.strip()
                if re.sub('[^a-zA-Z]', '', curr_transcript) == '':
                    curr_transcript = ''
                clip_transcript += " " + curr_transcript
                if end - curr_start > clip_duration:
                    clip_transcript = clip_transcript.replace("[", "").replace("]", "").replace(
                        "%", "").replace("(", "").replace(")", "").replace("-", " ").replace("=", "")
                    clip_transcript = unicodedata.normalize("NFKD", clip_transcript) \
                        .encode("ascii", "ignore")   \
                        .decode("ascii", "ignore")
                    clip_transcript = re.sub(r'\s+', ' ', clip_transcript).strip()
                    df_transcripts.loc[df_index] = [curr_start, end, clip_transcript, clip_duration, end - curr_start]
                    end_idx = int(end * sr)
                    curr_audio_data = audio_data[curr_start_idx: end_idx]
                    yt, index = librosa.effects.trim(curr_audio_data, top_db=10)
                    yt = curr_audio_data[max(index[0] - 40000, 0): min(index[1] + 40000, len(curr_audio_data))]
                    wav_file = FLAGS.data_dir + "/wav-files/" + file_prefix + str(file_index) + ".wav"
                    soundfile.write(wav_file, yt, sr)
                    file_index += 1
                    wav_filesize = os.path.getsize(wav_file)
                    data.append((os.path.abspath(wav_file), wav_filesize, clip_transcript))
                    df_index += 1
                    curr_start = end
                    curr_start_idx = end_idx
                    clip_transcript = ''
                    clip_duration = choice(clip_range)
            if clip_transcript != '':
                clip_transcript = clip_transcript.replace("[", "").replace("]", "").replace(
                    "%", "").replace("(", "").replace(")", "").replace("-", " ").replace("=", "")
                clip_transcript = unicodedata.normalize("NFKD", clip_transcript) \
                    .encode("ascii", "ignore")   \
                    .decode("ascii", "ignore")
                clip_transcript = re.sub(r'\s+', ' ', clip_transcript).strip()
                curr_audio_data = audio_data[curr_start_idx:]
                if len(curr_audio_data) > 0:
                    df_transcripts.loc[df_index] = [curr_start, end, clip_transcript, clip_duration, end - curr_start]
                    yt, index = librosa.effects.trim(curr_audio_data, top_db=10)
                    yt = curr_audio_data[max(index[0] - 40000, 0): min(index[1] + 40000, len(curr_audio_data))]
                    wav_file = FLAGS.data_dir + "/wav-files/" + file_prefix + str(file_index) + ".wav"
                    soundfile.write(wav_file, yt, sr)
                    file_index += 1
                    wav_filesize = os.path.getsize(wav_file)
                    data.append((os.path.abspath(wav_file), wav_filesize, clip_transcript))
                clip_transcript = ''
            df_transcripts.to_csv(FLAGS.data_dir + "/transcripts-" + str(idx) + ".csv", index=False)
            bar.update(1)
    pd.DataFrame(data=data, columns=["wav_filename", "wav_filesize", "transcript"]).to_csv(FLAGS.data_dir + "/train.csv", index=False)


if __name__ == "__main__":
    tf.compat.v1.app.flags.DEFINE_string("source_dir", "source dir", "source directory")
    tf.compat.v1.app.flags.DEFINE_string("data_dir", "data dir", "data directory")
    tf.compat.v1.app.flags.DEFINE_float("test_size", 0.02, "test size")
    tf.compat.v1.app.flags.DEFINE_integer("random_state", 1254, "random state")
    tf.compat.v1.app.run()

import tensorflow as tf
import librosa
import pandas as pd
import os
import unicodedata
import soundfile
import logging
import re
from tqdm import tqdm
from glob import glob
from sklearn.model_selection import train_test_split
from random import choice

logging.basicConfig(level=logging.NOTSET)
FLAGS = tf.compat.v1.app.flags.FLAGS

patterns = [r'(\.)+', r'(\([a-zA-Z0-9 _\(\)]+\))+', r'(<YWN)+', r'(?<=\s)=+\b', r'\b=+', r'\d(?=\w+)', r'(?<=\w)+\d', r'%(?=\w)+',
            r'\[(?=((\s|\'|\w|=|~|@|<|\-|\)|\()+))', r'(?<=(\s|\'|\w|=|~|@|>|\-|\)))+\]', r'(?<=\s)\-+', r'<[\w@%]+', r'[\w@%]+>', r'@+']


def main(_):
    source_dir = FLAGS.source_dir
    transcripts = glob(source_dir + "/transcripts/TRN/*.trn")
    clip_range = range(5, 10)
    data = []
    file_prefix = "sbc-"
    file_index = 1
    with tqdm(total=len(transcripts)) as bar:
        for idx, transcript_file in enumerate(transcripts):
            print(transcript_file)
            df = open(transcript_file, "r", encoding="utf-8")
            clip_duration = choice(clip_range)
            clip_transcript = ''
            curr_start = 0
            curr_start_idx = 0
            df_transcripts = pd.DataFrame(columns=["start", "end", "transcript", "selected clip_duration", "actual clip duration"])
            df_index = 0
            audio_file = source_dir + "/clips/" + transcript_file.split("/")[-1].split(".trn")[0] + ".wav"
            audio_data, sr = librosa.load(audio_file, sr=None)
            audio_duration = librosa.get_duration(audio_data)
            for i, line in enumerate(df):
                line = line.replace("\n", "")
                if line.strip() == "":
                    continue
                curr_transcript = line.split("\t")[-1]
                time_map = re.split(r'(\t|\s)+', line)
                start, end = time_map[0], time_map[2]
                if start == "":
                    start = time_map[2]
                    end = time_map[4]
                try:
                    start = float(start)
                    end = float(end)
                except:
                    print(f'{time_map}, {line}')
                if pd.isnull(curr_transcript):
                    curr_transcript = ''

                curr_transcript = re.sub(r'(' + "|".join(patterns) + r')', '', curr_transcript)
                curr_transcript = re.sub(r'((?<=\s)=+\b|\b=+|,|\?|~)', '', curr_transcript)
                curr_transcript = re.sub(r'\s+', ' ', curr_transcript)
                curr_transcript = curr_transcript.strip()
                if re.sub('[^a-zA-Z]', '', curr_transcript) == '':
                    curr_transcript = ''
                clip_transcript += " " + curr_transcript
                if end - curr_start > clip_duration:
                    clip_transcript = clip_transcript.replace("[", "").replace("]", "").replace(
                        "%", "").replace("(", "").replace(")", "").replace("-", " ").replace("=", "")
                    clip_transcript = unicodedata.normalize("NFKD", clip_transcript) \
                        .encode("ascii", "ignore")   \
                        .decode("ascii", "ignore")
                    clip_transcript = re.sub(r'\s+', ' ', clip_transcript).strip()
                    df_transcripts.loc[df_index] = [curr_start, end, clip_transcript, clip_duration, end - curr_start]
                    end_idx = int(end * sr)
                    curr_audio_data = audio_data[curr_start_idx: end_idx]
                    yt, index = librosa.effects.trim(curr_audio_data, top_db=10)
                    yt = curr_audio_data[max(index[0] - 40000, 0): min(index[1] + 40000, len(curr_audio_data))]
                    wav_file = FLAGS.data_dir + "/wav-files/" + file_prefix + str(file_index) + ".wav"
                    soundfile.write(wav_file, yt, sr)
                    file_index += 1
                    wav_filesize = os.path.getsize(wav_file)
                    data.append((os.path.abspath(wav_file), wav_filesize, clip_transcript))
                    df_index += 1
                    curr_start = end
                    curr_start_idx = end_idx
                    clip_transcript = ''
                    clip_duration = choice(clip_range)
            if clip_transcript != '':
                clip_transcript = clip_transcript.replace("[", "").replace("]", "").replace(
                    "%", "").replace("(", "").replace(")", "").replace("-", " ").replace("=", "")
                clip_transcript = unicodedata.normalize("NFKD", clip_transcript) \
                    .encode("ascii", "ignore")   \
                    .decode("ascii", "ignore")
                clip_transcript = re.sub(r'\s+', ' ', clip_transcript).strip()
                curr_audio_data = audio_data[curr_start_idx:]
                if len(curr_audio_data) > 0:
                    df_transcripts.loc[df_index] = [curr_start, end, clip_transcript, clip_duration, end - curr_start]
                    yt, index = librosa.effects.trim(curr_audio_data, top_db=10)
                    yt = curr_audio_data[max(index[0] - 40000, 0): min(index[1] + 40000, len(curr_audio_data))]
                    wav_file = FLAGS.data_dir + "/wav-files/" + file_prefix + str(file_index) + ".wav"
                    soundfile.write(wav_file, yt, sr)
                    file_index += 1
                    wav_filesize = os.path.getsize(wav_file)
                    data.append((os.path.abspath(wav_file), wav_filesize, clip_transcript))
                clip_transcript = ''
            df_transcripts.to_csv(FLAGS.data_dir + "/transcripts-" + str(idx) + ".csv", index=False)
            bar.update(1)
    pd.DataFrame(data=data, columns=["wav_filename", "wav_filesize", "transcript"]).to_csv(FLAGS.data_dir + "/train.csv", index=False)


if __name__ == "__main__":
    tf.compat.v1.app.flags.DEFINE_string("source_dir", "source dir", "source directory")
    tf.compat.v1.app.flags.DEFINE_string("data_dir", "data dir", "data directory")
    tf.compat.v1.app.flags.DEFINE_float("test_size", 0.02, "test size")
    tf.compat.v1.app.flags.DEFINE_integer("random_state", 1254, "random state")
    tf.compat.v1.app.run()
