#------------------------------------------------------------------------------
# Script
#------------------------------------------------------------------------------

import os
import glob
import json
import shutil
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--input_dir', default='./', type=str, help='Input directory')
parser.add_argument('--output_dir', default='./', 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]))

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

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

for base_name in ['models-1', 'models-2', 'models-3', 'models-4']:
    
    output_dir = os.path.join(args.output_dir, '%s-final' % base_name)
    print(output_dir)
    
    # Copy checkpoint dir for "aln" language as a base model.
    # We have only one best checkpoint always (no need to sort)
    aln_model_dir = glob.glob(os.path.join(args.input_dir, base_name, 'aln/checkpoint*'))[0]
    shutil.copytree(aln_model_dir, output_dir)
    
    # Copy adapters for all languages
    adapter_files = sorted(glob.glob(os.path.join(args.input_dir, base_name, '*/adapter.*.safetensors')))
    # print(len(adapter_files))
    
    for file in adapter_files:
        file_new = os.path.join(output_dir, os.path.basename(file))
        shutil.copy(file, file_new)
    
    # Join vocabs into single nested vocab and copy
    vocab_files = sorted(glob.glob(os.path.join(args.input_dir, base_name, '*/vocab.json')))
    # print(len(vocab_files))
    
    nested_vocab = {}
    for file in vocab_files:
        with open(file, 'rt', encoding='utf-8') as f:
            vocab = json.load(f)
        lang = list(vocab.keys())[0]
        nested_vocab[lang] = vocab[lang]
        
    # assert len(nested_vocab) == 26, 'Must be 26 langs'
    
    with open(os.path.join(output_dir, 'vocab.json'), 'wt', encoding='utf-8') as f:
        json.dump(nested_vocab, f)


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