#!/usr/bin/env python

from pangeamt_toolkit import Pangeanmt
from pangeamt_nlp import Pipeline
import os
import json
import shutil
import argparse
import subprocess

class Engine:
    def __init__(self, in_file, ref_file):
        # Load config
        with open('config.json', 'r') as file:
            self._config = json.loads(file.read())

        # Load src pipeline (used for training and decoding)
        self._src_pipeline = Pipeline(self._config['pipeline_config'],\
            self._config['src_lang'], self._config['tgt_lang'])

        # Load tgt pipeline (used only for training)
        self._tgt_pipeline = Pipeline(self._config['pipeline_config_tgt'],\
            self._config['tgt_lang'])

        self._in_file = in_file
        self._ref_file = ref_file

    def train(self):
        # Create translator object
        p = Pangeanmt('.')

        dir_to_in = os.path.join('data', self._in_file)
        dir_to_ref= os.path.join('data', self._ref_file)

        # Trains with the in_file and the ref_file
        with open(dir_to_in, 'r') as src_file:
            with open(dir_to_ref, 'r') as tgt_file:
                for seg in src_file:
                    try:
                        src = self._src_pipeline.preprocess_str(seg)
                        tgt = self._tgt_pipeline.preprocess_str(
                            tgt_file.readline()
                        )
                        p.train(src, tgt)
                    except:
                        print('Something went wrong.')
        return p

    def translate_file(self, p, output_file):
        dir_to_in = os.path.join('data', self._in_file)
        dir_to_out= os.path.join('data', output_file)

        # Translates the in_file and writes the postprocessed translation
        # in out_file
        with open(dir_to_in, 'r') as in_file:
            with open(dir_to_out, 'w+') as out_file:
                for seg in in_file:
                    self._src_pipeline.preprocess_str(seg)
                    translation = p.translate([seg])
                    tgt = (' ').join(translation[0].tgt)
                    tgt = self._src_pipeline.postprocess_str(tgt)
                    out_file.write(f'{tgt}\n')

    # Generates the config for the different learning rates
    def gen_config(self, alpha):
        lr = self._config['opts']['learning_rate']
        os.rename('config.json', f'{lr}_config.json')
        self._config['opts']['learning_rate'] = alpha
        with open('config.json', 'w+') as config_file:
            config_file.write(json.dumps(self._config))


# Create Parser
def _get_parser():
    parser = argparse.ArgumentParser(description='Try different learning rates'\
        ' to find the best.')
    parser.add_argument('dir', help="Path to extended model.")
    parser.add_argument('src_file', help='Which file to translate.')
    parser.add_argument('ref_file', help='Reference file when training.')
    parser.add_argument('learning_rate', help='Initial learning rate.')
    parser.add_argument('increment', help='Amount to increment each iteration.')
    parser.add_argument('max_learning_rate', help='Last learning rate to eval.')
    return parser

def main(args):
    # Change dir to current path
    os.chdir(args.dir)

    # Initialize Engine Object
    e = Engine(args.src_file, args.ref_file)

    # Create a translator to do a base translation
    p = Pangeanmt('.')
    e.translate_file(p, 'base_translation.txt')

    # Backup the original config
    shutil.copy('config.json', 'backup_config.json')

    # Loop with the different learning rate values
    alpha = args.learning_rate
    while alpha <= args.max_learning_rate:
        print(f'Training with learning rate = {alpha}')
        e.gen_config(alpha)

        # Train and return the translator with the trained model in memory
        p = e.train()
        print('Translating file.')
        output_file = f'{alpha}_{args.src_file}.txt'

        # Translate the file with the model in memory
        e.translate_file(p, output_file)
        params = ['pangeamt_multi_bleu.perl', args.ref_file]
        with open(output_file, 'r') as file:
            print(subprocess.check_output(params, stdin=file))

if __name__ == "__main__":
    parser = _get_parser()
    args = parser.parse_args()

    main(args)
