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

import os
import sys
import glob
import shutil
import subprocess
import numpy as np
import pandas as pd
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--input_dir', default='mdc_asr_shared_task_test_data', type=str, help='Directory with a test dataset')
parser.add_argument('--output_dir_single', default='submission_single_model', type=str, help='Output directory for a single model submission')
parser.add_argument('--output_dir_ensemble', default='submission_ensemble', type=str, help='Output directory for an ensemble submission')
parser.add_argument('--rover_path', default='./SCTK/bin/rover', type=str, help='Path to the ROVER binary')
args = parser.parse_args()
for a in [a for a in vars(args) if '__' not in a]: print('%-25s %s' % (a, vars(args)[a]))

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

def parse_rover_ctm(ctm_path):
    """
    Parse CTM ensemble file and convert to submission format

    Parameters:
    ctm_path : str
        Path to a CTM ensemble file

    Returns:
    result_df : pd.DataFrame
        DataFrame with final transcription strings
    """
    # Define CTM column names
    # CTM Format: ID channel start duration word confidence
    col_names = ['audio_file', 'channel', 'start', 'duration', 'word', 'conf']

    # Read the file
    df = pd.read_csv(ctm_path, sep=r'\s+', names=col_names, engine='python')

    # Sort by file ID and then by start time
    df = df.sort_values(by=['audio_file', 'start'])

    # Group by audio file and join words
    result_df = df.groupby('audio_file')['word'].apply(lambda x: ' '.join(str(w).strip() for w in x)).reset_index()
    # or with cleaning
    # result_df = df.groupby('audio_file')['word'].apply(lambda x: ' '.join(str(w).strip() for w in x if str(w) not in ['@', '*', '<unk>'])).reset_index()

    # Rename the column to 'sentence'
    result_df.rename(columns={'word': 'sentence'}, inplace=True)

    return result_df

#------------------------------------------------------------------------------
# Run ROVER ensembling
#------------------------------------------------------------------------------

os.makedirs(args.output_dir_single, exist_ok=True)
os.makedirs(args.output_dir_ensemble, exist_ok=True)

ctm_files = sorted(glob.glob('output_*_ctm/*/*.ctm'))
# print(len(ctm_files))

for file in ctm_files:
    shutil.copy(file, os.path.join(args.output_dir_ensemble, os.path.basename(file)))

tasks = ['multilingual-general', 'unseen-langs']
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']

langs = langs_general + langs_unseen

for lang in langs:

    res = subprocess.run(
        [
            args.rover_path,
            '-h', os.path.join(args.output_dir_ensemble, '%s_1.ctm' % lang), 'ctm',
            '-h', os.path.join(args.output_dir_ensemble, '%s_2.ctm' % lang), 'ctm',
            '-h', os.path.join(args.output_dir_ensemble, '%s_3.ctm' % lang), 'ctm',
            '-h', os.path.join(args.output_dir_ensemble, '%s_4.ctm' % lang), 'ctm',
            '-o', os.path.join(args.output_dir_ensemble, '%s_ens.ctm' % lang),
            '-m', 'avgconf',
            '-s',
            '-T',
        ],
        stdout=subprocess.DEVNULL,
        check=True,
    )

#------------------------------------------------------------------------------
# Parse ensemble CTM files and create TSV files
#------------------------------------------------------------------------------

for input_task in tasks:

    if input_task == 'multilingual-general':
        langs = langs_general
    else:
        langs = langs_unseen

    os.makedirs(os.path.join(args.output_dir_ensemble, input_task), exist_ok=True)

    for lang in langs:
        print('-'*50)
        print('Processing lang:', lang)

        #---- Load data
                
        input_file = os.path.join(args.input_dir, input_task, '%s.tsv' % lang)
        corpus_df = pd.read_csv(input_file, sep='\t')
        corpus_df['sentence'] = 'nanana'        
        print('Test size:', len(corpus_df))

        #---- Parse CTM ensemble file

        parsed_df = parse_rover_ctm(os.path.join(args.output_dir_ensemble, '%s_ens.ctm' % lang))

        #---- Save

        corpus_df = corpus_df.drop('sentence', axis=1)
        corpus_df = pd.merge(corpus_df, parsed_df, on='audio_file', how='left')
        
        out_file = os.path.join(args.output_dir_ensemble, input_task, os.path.basename(input_file))
        corpus_df[['audio_file', 'sentence']].to_csv(out_file, index=False, sep='\t')
        print('Saved:', out_file)

#------------------------------------------------------------------------------
# Create final submission from the best single model mms-1b-l1107
# We copy "multilingual-general" and "unseen-langs" from "output_2_submission", and "small-model" from "output_5_submission"
#------------------------------------------------------------------------------

shutil.copytree('output_2_submission/multilingual-general', os.path.join(args.output_dir_single, 'multilingual-general'))
shutil.copytree('output_2_submission/unseen-langs', os.path.join(args.output_dir_single, 'unseen-langs'))

os.makedirs(os.path.join(args.output_dir_single, 'small-model'), exist_ok=True)
files = sorted(glob.glob('output_5_submission/*/*.tsv'))
# print(len(files))
for file in files:
    shutil.copy(file, os.path.join(args.output_dir_single, 'small-model'))

#------------------------------------------------------------------------------
# Create final submission from an ensemble
# We already have "multilingual-general" and "unseen-langs" after CTM file parsing
# Now we copy "small-model" from "output_5_submission"
#------------------------------------------------------------------------------

os.makedirs(os.path.join(args.output_dir_ensemble, 'small-model'), exist_ok=True)
files = sorted(glob.glob('output_5_submission/*/*.tsv'))
# print(len(files))
for file in files:
    shutil.copy(file, os.path.join(args.output_dir_ensemble, 'small-model'))

# Transcription for "tob" language by the quantized model is the best, we copy it to "multilingual-general"
shutil.copy(os.path.join(args.output_dir_ensemble, 'small-model', 'tob.tsv'), os.path.join(args.output_dir_ensemble, 'multilingual-general'))

#------------------------------------------------------------------------------
# Create archives
#------------------------------------------------------------------------------

# Remove all .ctm files because we already created final .tsv files
files = sorted(glob.glob(os.path.join(args.output_dir_ensemble, '*.ctm')))
# print(len(files))
for file in files:
    os.remove(file)

# Create archives
zip_single_model = shutil.make_archive('submission_single_model', 'zip', root_dir=args.output_dir_single)
zip_ensemble = shutil.make_archive('submission_ensemble', 'zip', root_dir=args.output_dir_ensemble)

print('-'*50)
print('Saved as:', zip_single_model)
print('Saved as:', zip_ensemble)


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