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

import os
import re
import glob
import subprocess
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--input_dir', default='./kenlm_corpus', type=str, help='Input directory')
parser.add_argument('--output_dir', default='./kenlm_models_order_3', type=str, help='Output directory')
parser.add_argument('--bin_path', default='./kenlm_dist/build/bin/lmplz', type=str, help='Exec file')
parser.add_argument('--ngram_order', default=3, type=int, help='Ngram order')
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)

files = sorted(glob.glob(os.path.join(args.input_dir, '*.txt')))
print('N corpus files:', len(files))

for file in files:
    output_file = os.path.join(args.output_dir, os.path.basename(file).replace('.txt', '.arpa'))
    with open(file, 'rt', encoding='utf-8') as f_in:
        with open(output_file, 'wt', encoding='utf-8') as f_out:
            
            subprocess.run(
                [args.bin_path, '-o', str(args.ngram_order), '--discount_fallback'],
                stdin=f_in,
                stdout=f_out,
                check=True,
            )

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