#!/usr/bin/env python3 -u

import argparse
import fileinput
import logging
import os
import sys

from fairseq.models.transformer import TransformerModel


logging.getLogger().setLevel(logging.INFO)


def main():
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--en2fr', required=True,
                        help='path to en2fr model')
    parser.add_argument('--fr2en', required=True,
                        help='path to fr2en mixture of experts model')
    parser.add_argument('--user-dir',
                        help='path to fairseq examples/translation_moe/src directory')
    parser.add_argument('--num-experts', type=int, default=10,
                        help='(keep at 10 unless using a different model)')
    parser.add_argument('files', nargs='*', default=['-'],
                        help='input files to paraphrase; "-" for stdin')
    args = parser.parse_args()

    if args.user_dir is None:
        args.user_dir = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),  # examples/
            'translation_moe',
            'src',
        )
        if os.path.exists(args.user_dir):
            logging.info('found user_dir:' + args.user_dir)
        else:
            raise RuntimeError(
                'cannot find fairseq examples/translation_moe/src '
                '(tried looking here: {})'.format(args.user_dir)
            )

    logging.info('loading en2fr model from:' + args.en2fr)
    en2fr = TransformerModel.from_pretrained(
        model_name_or_path=args.en2fr,
        tokenizer='moses',
        bpe='sentencepiece',
    ).eval()

    logging.info('loading fr2en model from:' + args.fr2en)
    fr2en = TransformerModel.from_pretrained(
        model_name_or_path=args.fr2en,
        tokenizer='moses',
        bpe='sentencepiece',
        user_dir=args.user_dir,
        task='translation_moe',
    ).eval()

    def gen_paraphrases(en):
        fr = en2fr.translate(en)
        return [
            fr2en.translate(fr, inference_step_args={'expert': i})
            for i in range(args.num_experts)
        ]

    logging.info('Type the input sentence and press return:')
    for line in fileinput.input(args.files):
        line = line.strip()
        if len(line) == 0:
            continue
        for paraphrase in gen_paraphrases(line):
            print(paraphrase)


if __name__ == '__main__':
    main()
