import torch
import csv
import time
import argparse
from tqdm import tqdm
import editdistance
import unicodedata
from pororo import Pororo

def save_csv(hypothesis, ground_truth, path="result.csv"):
    f = open(path, "w", newline="")
    wr = csv.writer(f, delimiter=' ')

    for i in range(len(ground_truth)):
        wr.writerows([ground_truth[i], hypothesis[i]])

    f.close()

def load_script(temp_path):
    with open(temp_path, 'r') as f:
        lines = list()
        for line in f.readlines():
            lines.append(line.replace('\n', ''))
        return lines

def get_scores(refs, hyps):
    errs_wer = 0
    lengths_wer = 0

    errs_cer = 0
    lengths_cer = 0
    
    wer = None
    cer = None
    
    ## normalize(한글 NFKD 일 경우 글자 갯수가 다르게 측정되는 걸 방지하기 위하여)
    for tgt_word, hyp_word in tqdm(zip(refs, hyps), total=len(refs)):
        hyp_word = unicodedata.normalize("NFC", hyp_word)
        tgt_word = unicodedata.normalize("NFC", tgt_word)

        err_w, length_w = editdistance.eval(hyp_word.split(), tgt_word.split()), len(tgt_word.split())
        err_c, length_c = editdistance.eval(hyp_word, tgt_word), len(tgt_word)

        errs_wer += err_w
        lengths_wer += length_w
        errs_cer += err_c
        lengths_cer += length_c
        
    
    if lengths_wer > 0:
        wer = errs_wer * 100.0 / lengths_wer


    if lengths_cer > 0:
        cer = errs_cer * 100.0 / lengths_cer

    return cer, wer


# def infer(wav_path:str, asr):
#     start = time.time()
#     results = asr(wav_path)['results']
#     if len(results) != 0:
#         hypothesis = results[0]['speech']
#     return hypothesis

def infer(wav_path:str):
    start = time.time()
    results = asr(wav_path)['results'][0]['speech']
    return results




if __name__ == "__main__":
    EVAL_CLEAN = '/data/STT_data/KsponSpeech-mini-16000/eval_clean_transcript.txt'
    EVAL_OTHER = '/data/STT_data/KsponSpeech-mini-16000/eval_other_transcript.txt'
    TEST = '/data/STT_data/KsponSpeech-mini/test_transcript.txt'

    EVAL_CLEAN_WAV = '/data/STT_data/KsponSpeech-mini-16000/eval_clean_wav/'
    EVAL_OTHER_WAV = '/data/STT_data/KsponSpeech-mini-16000/eval_other_wav/'
    TEST_WAV = '/data/STT_data/KsponSpeech-mini/test_wav/'

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("--split", type=str, default="eval_clean")
    config = parser.parse_args()
    print(config)
    if config.split == "eval_clean":
        script = EVAL_CLEAN
        wav_dir = EVAL_CLEAN_WAV

    elif config.split == "eval_other":
        script = EVAL_OTHER
        wav_dir = EVAL_OTHER_WAV
    
    elif config.split == "test":
        script = TEST
        wav_dir = TEST_WAV

    print('load the model')
    asr = Pororo(task='asr', lang='ko')

    ground_truth = []
    hypothesis = []

    with open(script) as f:
        lines = f.readlines()
        start = time.time()
        for line in lines:
            line = line.strip()
            wav = wav_dir + line[:18] + ".wav"
            ground_truth.append(line[19:])
            hypothesis.append(infer(wav))

    print("time: ", time.time() - start)
    
    f1 = open('ground_truth.txt', 'w')  
    for i in range(len(ground_truth)):
        f1.write(ground_truth[i]+'\n')
    f1.close()

    f2 = open('hypothesis.txt', 'w')  
    for i in range(len(hypothesis)):
        f2.write(hypothesis[i]+'\n')
    f2.close()

    print("complete")

    cer, wer= get_scores(ground_truth, hypothesis)
    print('cer: ', cer)
    print('wer: ', wer)


    ground_truth = load_script('./ground_truth.txt')
    hypothesis = load_script('./hypothesis.txt')

    a,b = get_scores(ground_truth, hypothesis)
    print(a, b)