#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# author : Santosh Kesiraju
# e-mail : kcraj2[AT]gmail[DOT]com, kesiraju[AT]fit[DOT]vutbr[DOT]cz
# Date created : 19 Jul 2022
# Last modified : 19 Jul 2022

"""
Parse log file and get best epoch based on validation score
detokenize the reference and hypothesis
use sacre_bleu.pl for BLEU
"""

import os
import sys
import json
import subprocess
import argparse
from typing import List
import numpy as np
from sacrebleu.metrics import BLEU, CHRF, TER


MOSES = os.environ.get("MOSES", None)
assert MOSES != None, "Please set the env variable MOSES to the mosesdecoder root dir."
DETOK = os.path.join(MOSES, "scripts/tokenizer/detokenizer.perl")
assert os.path.exists(
    DETOK
), "Cannot find detokenzier.perl at {:s}. Make sure your MOSES root dir path is correct.".format(
    DETOK
)


def read_simple_flist(fname, pre="", sfx=""):
    """read simple file line by line into list.
    Append the prefix and suffix for every line if given

    Args:
    fname (str): full path to file
    pre (str): string for prefix
    sfx (str): string for suffix

    Returns:
    list
    """

    flist = []
    with open(fname, "r", encoding="utf-8") as fpr:
        flist = [line.strip() for line in fpr if line.strip()]

    if flist[-1].strip() == "":
        flist = flist[:-1]

    if pre != "" or sfx != "":
        flist_new = [pre + f + sfx for f in flist]
        flist = flist_new

    return flist


def write_simple_flist(some_list, out_fname):
    """Write the elements in the list line by line
    in the given out file

    Parameters:
    -----------
    some_list (list): list of elements
    out_fname (str): output file name

    """

    with open(out_fname, "w", encoding="utf-8") as fpw:
        fpw.write("\n".join(some_list))


def parse_log_file_and_load_results(input_log_file) -> List:
    """Parse log or json file and load results into list of dict"""

    results = []
    dump_path = ""
    hyp_path = ""
    with open(input_log_file, "r") as fpr:
        for line in fpr:
            line = line.strip()
            if line.startswith("INFO"):
                parts = line.split(" - ")
                if parts[-1].startswith("__log__"):
                    json_str = parts[-1].split("__log__:")[-1]
                    results.append(json.loads(json_str))
            elif line.startswith("dump_path"):
                dump_path = line.split(":")[-1].strip()
            else:
                continue

    if dump_path:
        hyp_path = os.path.join(dump_path, "hypotheses")
        assert os.path.isdir(hyp_path), f"Could not find hypotheses dir {hyp_path}"
    else:
        print("Could not find dump_path in the log file", file=sys.stderr)
        sys.exit()

    print("- Loaded results:", len(results), len(results[0]))

    return results, hyp_path


def get_best_val_epoch(results, args):
    """Get the epoch number that gives the best val score (args.key)"""

    epochs = []
    metric_name = ""
    keys = []
    scores = {args.dev_set: {}}
    for tst in args.test_sets:
        scores[tst] = {}

    epochs = []
    scores = {"valid": {}, "test": {}}
    metric_name = ""
    keys = []

    for k in results[0]:
        if k == "epoch":
            continue
        else:
            if k.startswith("valid"):
                key = k.split("_", 1)[-1]
                keys.append(key)

    print("- Metrics (keys):", keys)

    for entry in results:
        for k, v in entry.items():
            if k == "epoch":
                epochs.append(entry["epoch"])
            else:
                set_name, metric_name = k.split("_", 1)
                key_ix = keys.index(metric_name)
                try:
                    scores[set_name][key_ix].append(v)
                except KeyError:
                    scores[set_name][key_ix] = [v]

    print(keys)
    keyix = keys.index(args.lang_pair + "_mt_" + args.key)
    metric_name = keys[keyix]
    print(keyix, metric_name)

    print("- Using metric:", metric_name)
    best_ix = np.argmax(scores["valid"][keyix])

    print("==== Scores before detokenization ====")
    for set_name in scores:
        print("{:10s} {:4.2f}".format(set_name, scores[set_name][keyix][best_ix]))
    print("=" * 38)

    return scores, best_ix


def run_detokenizer_bleu(epoch_ix, hyp_path, lang_pair, set_names, res_dir):
    """Run detokenizer"""

    src_lang, tar_lang = lang_pair.split("-")

    for base in ["ref", f"hyp{epoch_ix}"]:
        for set_name in set_names:

            ref_txt = os.path.join(hyp_path, f"{base}.{lang_pair}.{set_name}.txt")
            ref_detok = os.path.join(
                hyp_path, f"{base}.{lang_pair}.{set_name}.txt.detok"
            )
            assert os.path.exists(ref_txt), "- Cannot find text file at {:s}".format(
                ref_txt
            )

            if os.path.exists(ref_detok):
                print("- Detokenized file already exists. Skipping..", ref_detok)
                continue

            fpr = open(ref_txt, "r")
            fpw = open(ref_detok, "w")
            print("- Detokenizing", ref_txt)
            proc = subprocess.Popen(
                [DETOK, "-l", src_lang, "-q"], stdin=fpr, stdout=fpw
            )
            try:
                proc.wait(timeout=120)
            except subprocess.TimeoutExpired:
                print("- Timeout expired detokenizer.perl")
            finally:
                fpr.close()
                fpw.close()

    for set_name in set_names:

        ref_detok = os.path.join(hyp_path, f"ref.{lang_pair}.{set_name}.txt.detok")
        hyp_detok = os.path.join(
            hyp_path, f"hyp{epoch_ix}.{lang_pair}.{set_name}.txt.detok"
        )

        res_file = os.path.join(res_dir, f"{set_name}.txt")

        ref = read_simple_flist(ref_detok)
        hyp = read_simple_flist(hyp_detok)

        bleu = BLEU()
        bleu_str = bleu.corpus_score(hyp, [ref]).format()
        bleu_str += "\n" + bleu.get_signature().format()

        chrf = CHRF()
        chrf_str = chrf.corpus_score(hyp, [ref]).format()
        chrf_str += "\n" + chrf.get_signature().format()

        ter = TER()
        ter_str = ter.corpus_score(hyp, [ref]).format()
        ter_str += "\n" + ter.get_signature().format()

        print(bleu_str, chrf_str, ter_str, sep="\n")

        write_simple_flist([bleu_str, chrf_str, ter_str], res_file)
        print("- Results saved in", res_file)


def main():
    """main method"""

    args = parse_arguments()

    if not args.results_dir:
        base = os.path.basename(args.input_log_file).rsplit(".", 1)[0]
        args.results_dir = os.path.join(os.path.dirname(args.input_log_file), f"results_{base}/")

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

    results, hyp_path = parse_log_file_and_load_results(args.input_log_file)

    scores, epoch_ix = get_best_val_epoch(results, args)

    run_detokenizer_bleu(
        epoch_ix, hyp_path, args.lang_pair, list(scores.keys()), args.results_dir
    )


def parse_arguments():
    """parse command line arguments"""

    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("input_log_file", help="path to train log file")
    parser.add_argument(
        "lang_pair", type=str, help="language pair: en-de if going from en -> de"
    )
    parser.add_argument("-dev_set", type=str, default="val", help="name of dev set")
    parser.add_argument(
        "-test_sets", type=str, nargs="+", help="list of test sets", required=True
    )
    parser.add_argument(
        "-key",
        choices=["ppl", "acc", "bleu"],
        type=str,
        default="bleu",
        help="metric to choose the best validation model/output",
    )
    parser.add_argument(
        "-results_dir", type=str, default="",
        help="path to output dir to save results. defauls to results/ in same dir as log file")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    main()
