#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

import os
import re
import pandas as pd
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--input_dir', default='./', type=str, help='Input directory')
parser.add_argument('--output_dir', default='./kenlm_corpus', type=str, help='Output directory')
args = parser.parse_args()
for a in [a for a in vars(args) if '__' not in a]: print('%-25s %s' % (a, vars(args)[a]))

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

bracketed = re.compile(r"\[[^\]]+\]")
unintell_paren = re.compile(r"\(\?+\)")
repl_punc = re.compile('[,?¿¡!";:]+')
multispace = re.compile("  +")

def clean(t):
    """
    Official cleaning function
    """
    t = re.sub(bracketed, " ", t)
    t = re.sub(unintell_paren, " ", t)
    t = t.replace(" ... ", " ")
    t = t.replace("#x27;", "'")
    t = re.sub(repl_punc, " ", t)
    t = t.replace("...", "!ELLIPSIS!").replace(".", " ").replace("!ELLIPSIS!", "...")
    t = re.sub(multispace, " ", t)
    return t

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

os.makedirs(args.output_dir, exist_ok=True)

langs_general = ['aln', 'bew', 'bxk', 'cgg', 'el-CY', 'hch', 
                 'kcn', 'koo', 'led', 'lke', 'lth', 'meh', 'mmc', 
                 'pne', 'ruc', 'rwm', 'sco', 'tob', 'top', 'ttj', 'ukv']
langs_unseen = ['ady', 'bas', 'kbd', 'qxp', 'ush']

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

for lang in langs_general:
    corpus_df = pd.read_csv(os.path.join(args.input_dir, 'mcv-sps-st-09-2025/sps-corpus-1.0-2025-09-05-%s/ss-corpus-%s.tsv' % (lang, lang)), sep='\t')
    print('Lang: %s    Size: %d' % (lang, len(corpus_df)))

    selector_no_trans = corpus_df['transcription'].isnull()
    corpus_df = corpus_df[~selector_no_trans]
    selector_zero_len_trans = corpus_df['transcription'].map(len) == 0
    corpus_df = corpus_df[~selector_zero_len_trans]
    corpus_df['transcription'] = corpus_df['transcription'].map(clean)

    with open(os.path.join(args.output_dir, '%s.txt' % lang), 'wt', encoding='utf-8') as f:
        for line in corpus_df['transcription'].values:
            f.write(line + '\n')

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

for lang in langs_unseen:
    train_df = pd.read_csv(os.path.join(args.input_dir, 'cv-corpus-23.0-2025-09-05/%s/train.tsv' % lang), sep='\t')
    dev_df = pd.read_csv(os.path.join(args.input_dir, 'cv-corpus-23.0-2025-09-05/%s/dev.tsv' % lang), sep='\t')    
    test_df = pd.read_csv(os.path.join(args.input_dir, 'cv-corpus-23.0-2025-09-05/%s/test.tsv' % lang), sep='\t')  
    corpus_df = pd.concat([train_df, dev_df, test_df])
    corpus_df['transcription'] = corpus_df['sentence']  # just for compat
    print('Lang: %s    Size: %d' % (lang, len(corpus_df)))

    selector_no_trans = corpus_df['transcription'].isnull()
    corpus_df = corpus_df[~selector_no_trans]
    selector_zero_len_trans = corpus_df['transcription'].map(len) == 0
    corpus_df = corpus_df[~selector_zero_len_trans]
    corpus_df['transcription'] = corpus_df['transcription'].map(clean)

    with open(os.path.join(args.output_dir, '%s.txt' % lang), 'wt', encoding='utf-8') as f:
        for line in corpus_df['transcription'].values:
            f.write(line + '\n')

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------