#!/usr/bin/env python
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import argparse
import logging
import os
import re
import shlex
import subprocess
import sys

FORMAT = '%(levelname)s: %(message)s'
logging.basicConfig(level=0, format=FORMAT)


def handle_args(argv):
    program = argv[0]
    arguments = argv[1:]

    class MyParser(argparse.ArgumentParser):
        def error(self, message=None):
            if message:
                logging.error('%s' % message)
            self.print_help()
            sys.exit(2)

    parser = MyParser(
            prog=program,
            description='Merge multiple language models into a single one.',
            epilog='Interpolation of the input models is based on a '
                   'perplexity measurement of each model against the '
                   'development text.'
    )
    parser.add_argument(
            'input_lms',
            nargs='+',
            help='paths to language models to be merged'
    )
    parser.add_argument(
            'dev_text',
            help='path to a file that will be used for calculating '
                 'interpolation weights for the language models to be merged'
    )
    parser.add_argument(
            'merged_lm_path',
            help='path to where the output merged LM file should be written'
    )
    parser.add_argument(
            '--temp-dir', default=os.path.join(os.getcwd(), 'merge-lms-tmp'),
            help='path to the directory where perplexity calculations will '
                 'be stored. "./%(default)s/" is the default location. The '
                 'temp dir is not automatically deleted.'
    )
    # Determine the path to the srilm tool's binaries
    try:
        srilm_env_var_val = os.environ['SRILM']
    except KeyError:
        srilm_env_var_val = 'not set'
    parser.add_argument(
            '--srilm-bin',
            default=os.path.join(srilm_env_var_val, 'bin', 'i686-m64'),
            help='path to where the srilm tool\'s binaries have been '
                 'compiled. By default this is "$SRILM/bin/i686-m64". '
                 'Use this option if your srilm binaries are compiled to a '
                 'different location. The SRILM environment variable is '
                 'currently set to "%s" in your system. '
                 % (srilm_env_var_val)
    )

    args = parser.parse_args(arguments)

    fail = False
    # Assert that there are at least 2 input LM files.
    if len(args.input_lms) < 2:
        fail = 2
        logging.error('2 or more input LM files are required.')

    # Assert that all the input LMs exist.
    for lm_path in args.input_lms:
        if not os.path.exists(lm_path):
            fail = 2
            logging.error("The input LM '%s' was not found." % (lm_path))

    # Assert that the dev text exists.
    if not os.path.exists(args.dev_text):
        fail = 2
        logging.error("The dev text '%s' was not found." % (args.dev_text))

    # If merged-lm-path already exists, warn that it is about to be
    # overwritten.
    if os.path.exists(args.merged_lm_path):
        logging.warn(
                'The merged LM output destination %s already exists and will '
                'be overwritten.' % args.merged_lm_path
        )

    # Handle the temp dir.
    if not os.path.isdir(args.temp_dir):
        if os.path.exists(args.temp_dir):
            fail = 2
            logging.error("%s is not a directory." % args.temp_dir)
        else:
            logging.info("Making temporary directory %s" % args.temp_dir)
            os.makedirs(args.temp_dir)

    # Assert that srilm_bin has a variable
    if args.srilm_bin is 'not set':
        parser.error(
                message="No value was provided for '--srilm-bin', and the "
                        "SRILM environment variable is not set."
        )

    # Assert that the srilm tools are located in the location specified.
    for tool in ['ngram', 'compute-best-mix']:
        if not os.path.exists(os.path.join(args.srilm_bin, tool)):
            fail = 2
            logging.error("The SRILM tool '%s' was not found in %s."
                          % (tool, args.srilm_bin))

    if fail:
        sys.exit(fail)

    return args


def exec_shell(cmd_str):
    logging.info(cmd_str)
    p = subprocess.Popen(
            shlex.split(cmd_str),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
    )
    try:
        stdoutdata, stderrdata = p.communicate()
        logging.info(stdoutdata)
        logging.info(stderrdata)
    except OSError as xcpt:
        logging.info(stdoutdata)
        logging.error(stderrdata)
        logging.error(xcpt.args[1])
        sys.exit(xcpt.args[0])
    return stdoutdata, stderrdata


def best_mix(args):
    """
    Calculate the perplexity of each language model against all the dev text.
    Write the perplexity calculations to files. Use the compute-best-mix script
    to determine interpolation weights. Returns the interpolation weights.
    """
    # Calculate perplexities
    ppl_destinations = [
            os.path.join(args.temp_dir, 'lm-{}.ppl'.format(i))
            for i in range(len(args.input_lms))
    ]
    for i, lm in enumerate(args.input_lms):
        cmd = ' '.join([
                os.path.join(args.srilm_bin, 'ngram'),
                '-debug 2 -order 5 -unk',
                '-lm', lm,
                '-ppl', args.dev_text,
        ])
        stdoutdata, stderrdata = exec_shell(cmd)
        with open(ppl_destinations[i], 'w') as fh:
            fh.write(stdoutdata)
    # Compute best mix
    cmd = '{} {}'.format(
            os.path.join(args.srilm_bin, 'compute-best-mix'),
            ' '.join(ppl_destinations),
    )
    best_mix_result, stderrdata = exec_shell(cmd)
    best_mix_ppl_path = os.path.join(args.temp_dir, 'best-mix.ppl')
    with open(best_mix_ppl_path, 'w') as fh:
        logging.info('Writing %s' % best_mix_ppl_path)
        fh.write(best_mix_result)
    weights = parse_lambdas(best_mix_result)
    if len(weights) != len(ppl_destinations):
        sys.exit('The number of weights found in %s is inconsistent with the '
                 'number of .ppl files.')
    return weights


def parse_lambdas(data):
    m = re.search(r"\((.*)\)", data)
    if not m:
        logging.error('The lambda values were not found.')
        sys.exit(-1)
    return m.group(1).split()


def construct_merge_cmd(weights, args):
    result = """
            {0}/ngram -order 5 -unk
            -lm     {1[0]} -lambda {2[0]}
            -mix-lm {1[1]}
    """.format(args.srilm_bin, args.input_lms, weights)
    idx = 2
    while idx < len(weights):
        result += """
                -mix-lm{0} {1} -mix-lambda{0} {2}
        """.format(idx, args.input_lms[idx], weights[idx])
        idx += 1
    return ' '.join(result.split() + ['-write-lm', args.merged_lm_path])


def merge_lms(weights, args):
    cmd = construct_merge_cmd(weights, args)
    logging.info('Merging LMs int %s' % args.merged_lm_path)
    exec_shell(cmd)


if __name__ == '__main__':
    args = handle_args(sys.argv)
    weights = best_mix(args)
    merge_lms(weights, args)