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

# Santosh

"""
Plot results (ppl, acc, bleu) from json file
"""

import argparse
import json
import numpy as np


def parse_file_and_load_results(args):
    """Parse log or json file and load results into dict"""

    if args.in_json_file.endswith(".log"):
        print(
            "- Input file is a log file. Will extract results and save it in a json file"
        )
        res_file = args.in_json_file.replace(".log", ".json")

        json_lines = []
        with open(args.in_json_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]
                        json_lines.append(json_str)
        print("- Found", len(json_lines), "json lines.")
        with open(res_file, "w") as fpw:
            fpw.write("\n".join(json_lines))

        print("- Json results file saved:", res_file)
        args.in_json_file = res_file

    results = []
    try:
        with open(args.in_json_file, "r", encoding="utf-8") as fpr:
            data = fpr.readlines()
            for row in data:
                results.append(json.loads(row))
    except json.decoder.JSONDecodeError:
        for row in data:
            results.append(row.strip())

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


def get_best_val_epoch(results, args):
    """Get best val epoch and corresponding scores"""

    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]

    if args.keyix == -1:
        args.keyix = len(keys) - 1
    metric_name = keys[args.keyix]

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

    print(
        f"- Best valid {keys[-1]} at epoch:",
        epochs[best_ix],
        "valid:",
        best_val,
        "test:",
        best_test,
    )

    return epochs, scores, keys, metric_name, [best_ix, best_val, best_test]


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

    args = parse_arguments()

    results = parse_file_and_load_results(args)
    epochs, scores, keys, metric_name, best = get_best_val_epoch(results, args)
    best_ix, best_valid, best_test = best

    if args.show or args.save:
        import matplotlib.pyplot as plt

        nrows = 2
        ncols = len(keys) // 2
        plt.style.use("presentation")
        plt.rc("text", usetex=True)
        if ncols == 1:
            nrows += 1
        fig, axs = plt.subplots(
            num=1, nrows=nrows, ncols=ncols, sharex=True, figsize=(14, 8)
        )
        axs = axs.flatten()

        for i, key in enumerate(keys):

            axs[i].plot(epochs, scores["valid"][i], "-", color="C0", alpha=0.5)
            axs[i].plot(epochs, scores["test"][i], "-", color="C1", alpha=0.5)
            axs[i].set_title(key.replace("_", "-"))
            axs[i].set_xlabel("epochs")
            axs[i].grid(alpha=0.5, linestyle="--")

        axs[-1].plot(
            epochs[best_ix],
            scores["valid"][args.keyix][best_ix],
            "*",
            color="C0",
            label="valid {:s} {:4.2f}".format(
                metric_name.replace("_", "-"), best_valid
            ),
        )
        axs[-1].plot(
            epochs[best_ix],
            scores["test"][args.keyix][best_ix],
            "*",
            color="C1",
            label="test {:s} {:4.2f}".format(metric_name.replace("_", "-"), best_test),
        )

        axs[-1].legend(loc="best")

        plt.tight_layout()

    if args.show:
        plt.show()

    if args.save:
        plt.savefig(args.save, dpi=300)
        print("- Saved", args.save)


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

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "in_json_file",
        help="path to json file with log or results"
    )
    parser.add_argument(
        "-keyix", type=int, default=-1, help="index of eval metric (key ix)"
    )
    parser.add_argument(
        "-save", default="", help="path to file with ext, to save the image"
    )
    parser.add_argument("--show", action="store_true", help="show the plot")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    main()
