# --------------------------------------------------------------------------------------------------
# SPDX-License-Identifier: Apache-2.0
# Copyright (C) 2024 Jayesh Badwaik <j.badwaik@fz-juelich.de>
# --------------------------------------------------------------------------------------------------

import jureap.detail
import shutil
import itertools
import os
import tempfile
import subprocess
import json
import sys
import matplotlib.pyplot
import math
import numpy
import matplotlib.ticker


def compute_data_list(input_dir):
    data_name_list = []
    for file in os.listdir(input_dir):
        if file.endswith(".csv"):
            data_name_list.append(file.split(".")[0])

    return data_name_list


def extract_nodes_and_runtime(input_file):
    data = jureap.detail.csv_file_to_array(input_file)
    nodes = []
    runtime = []
    for row in data:
        nodes.append(int(row["nodes"]))
        runtime.append(float(row["runtime"]))

    runtime = [x for _, x in sorted(zip(nodes, runtime))]
    nodes = sorted(nodes)
    ideal_scaling_array = [runtime[0] / n * nodes[0] for n in nodes]
    low_scaling_array = [1.25 * i for i in ideal_scaling_array]

    return {
            "nodes": nodes,
            "runtime": runtime,
            "ideal_scaling": ideal_scaling_array,
            "low_scaling": low_scaling_array,
            }


def extract_data(input_dir):
    workload_factor = read_workload_factor(input_dir)
    print(workload_factor)
    min_workload_factor = min(workload_factor.values())
    for key in workload_factor:
        workload_factor[key] = workload_factor[key] / min_workload_factor

    plot_data = []
    for file in os.listdir(input_dir):
        if file.endswith(".csv"):
            data = extract_nodes_and_runtime(os.path.join(input_dir, file))
            name = file.rsplit(".", 1)[0]
            plot_data.append(
                    {
                        "name": file.rsplit(".", 1)[0],
                        "data": data,
                        "workload_factor": workload_factor[name],
                        }
                    )

    return plot_data


def read_workload_factor(input_dir):
    workload_factor = {}
    for file in os.listdir(input_dir):
        if file.endswith(".csv"):
            basename = file.rsplit(".", 1)[0]
            workload_filename = file.rsplit(".", 1)[0] + ".workload"
            with open(os.path.join(input_dir, workload_filename), "r") as f:
                workload_factor[basename] = float(f.read().strip())

    return workload_factor


def prepare_output_dir(input_dir, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)
    shutil.copytree(input_dir, os.path.join(output_dir, "input"))


def compute_plot_info(plot_data):
    plot_info = {}
    node_range = [sys.maxsize, 0]
    runtime_range = [sys.float_info.max, 0]

    min_plot_name = plot_data[0]["name"]
    min_node = sys.maxsize

    node_array = []

    for keydata in plot_data:
        data = keydata["data"]
        name = keydata["name"]
        if min(data["nodes"]) < min_node:
            min_plot_name = name
            min_node = min(data["nodes"])
        node_array = node_array + data["nodes"]
        node_range[0] = min(node_range[0], min(data["nodes"]))
        node_range[1] = max(node_range[1], max(data["nodes"]))
        runtime_range[0] = min(runtime_range[0], min(data["runtime"]))
        runtime_range[1] = max(runtime_range[1], max(data["runtime"]))

    for plot in plot_data:
        if plot["name"] == min_plot_name:
            base_expected_runtime = plot["workload_factor"] * plot["data"]["runtime"][0]

    expected_runtime_array = {}
    for plot in plot_data:
        name = plot["name"]
        data = plot["data"]
        expected_runtime_array[name] = (
                base_expected_runtime * plot["workload_factor"] / data["nodes"][0] * min_node
                )

    node_range = [node_range[0] * 0.75, node_range[1] * 1.25]
    runtime_range = [runtime_range[0] * 0.75, runtime_range[1] * 1.12]

    node_array = sorted(list(set(node_array)))
    reduced_node_array = [node_array[0]]
    last_position = 0
    for i in range(len(node_array) - 1):
        if node_array[i + 1] / reduced_node_array[-1] > 1.5:
            reduced_node_array.append(node_array[i + 1])

    if max(node_array) not in reduced_node_array:
        reduced_node_array.append(max(node_array))

    if reduced_node_array[-2] / reduced_node_array[-1] > 0.66:
        reduced_node_array.pop(-2)
    node_array = reduced_node_array

    plot_info["xticklabels"] = node_array
    yticklabel_range = [math.log(runtime_range[0]), math.log(runtime_range[1])]
    plot_info["yticklabels"] = numpy.logspace(
            yticklabel_range[0], yticklabel_range[1], num=6, base=math.e
            )
    plot_info["yticklabels"] = [int(p) if p > 1 else p for p in plot_info["yticklabels"]]
    plot_info["min_plot_name"] = min_plot_name
    plot_info["range_limits"] = {"node": node_range, "runtime": runtime_range}

    plot_info["min_plot_name"] = min_plot_name
    plot_info["expected_runtime_array"] = expected_runtime_array
    return plot_info


def generate_plot_file(system_name, output_dir, plot_data, skip_weak_scaling, legend_position):
    plot = matplotlib.pyplot.figure()
    marker = itertools.cycle(("o", "^", "s", "X"))

    ax = plot.add_subplot(111, label="1")
    ax.ticklabel_format(useOffset=False)

    cmap = matplotlib.colormaps["tab10"]
    colorgen = iter([cmap(i) for i in range(len(plot_data))])
    plot_info = compute_plot_info(plot_data)

    error_yaxis = []
    xaxis_point = []
    yaxis_point = []

    color_history = []

    labeled_interval = False

    ax.fill_between(
            [0, 0],
            [0, 0],
            [0, 0],
            label="80% parallel efficiency intervals",
            alpha=0.2,
            ls="--",
            color="grey",
            )
    for data in plot_data:
        color = next(colorgen)
        color_history.append(color)
        ax.plot(
                data["data"]["nodes"],
                data["data"]["runtime"],
                marker=next(marker),
                color=color,
                )
        if not labeled_interval:
            labeled_interval = True
        else:
            label = None
        ax.fill_between(
                data["data"]["nodes"],
                data["data"]["ideal_scaling"],
                data["data"]["low_scaling"],
                alpha=0.2,
                ls="--",
                color=color,
                )
        if data["name"] != plot_info["min_plot_name"]:
            expected_runtime = plot_info["expected_runtime_array"][data["name"]]
            error_yaxis_positive = 0.25 * expected_runtime
            error_yaxis_negative = 0
            error_yaxis.append([error_yaxis_negative, error_yaxis_positive])
            xaxis_point.append(data["data"]["nodes"][0])
            yaxis_point.append(expected_runtime)

            error_yaxis_transpose = [[], []]
            for point in error_yaxis:
                for i in range(2):
                    error_yaxis_transpose[i].append(point[i])

            minimum_error_bar_ymin = min(
                    [yxp - yerr for yxp, yerr in zip(yaxis_point, error_yaxis_transpose[0])]
                    )
            max_error_bar_ymax = max(
                    [yxp + yerr for yxp, yerr in zip(yaxis_point, error_yaxis_transpose[1])]
                    )

            if minimum_error_bar_ymin < plot_info["range_limits"]["runtime"][0]:
                plot_info["range_limits"]["runtime"][0] = minimum_error_bar_ymin
            if max_error_bar_ymax > plot_info["range_limits"]["runtime"][1]:
                plot_info["range_limits"]["runtime"][1] = max_error_bar_ymax

            error_bar_ymin = min(min(yaxis_point), min(error_yaxis_transpose[0]))

    if len(plot_data) > 1:
        if not skip_weak_scaling:
            eb = ax.errorbar(
                    xaxis_point,
                    yaxis_point,
                    label="80% weak scaling efficiency intervals\n(estimated relative to first red datum)",
                    yerr=error_yaxis_transpose,
                    fmt="none",
                    capsize=4.0,
                    ls="-",
                    alpha=0.7,
                    )
            eb[-1][0].set_linestyle("dashed")

    ax.set_xscale("log")
    ax.set_yscale("log")

    ax.set_xlim(plot_info["range_limits"]["node"])
    ax.set_ylim(plot_info["range_limits"]["runtime"])

    ax.set_xticks(plot_info["xticklabels"])
    ax.set_yticks(plot_info["yticklabels"])
    ax.minorticks_off()

    ax_t = ax.secondary_xaxis("top")
    ax_t.tick_params(axis="x", labelrotation=90)
    ax_t.minorticks_off()
    ax_t.set_xticks(plot_info["xticklabels"])
    ax_t.set_xticklabels([4 * i for i in plot_info["xticklabels"]])
    ax_t.set_xlabel("Number of NVIDIA A100 GPUs (log-scale)")

    ax.set_xticklabels(plot_info["xticklabels"])
    ax.set_yticklabels(plot_info["yticklabels"])
    if plot_info["range_limits"]["runtime"][1] < 1 or plot_info["range_limits"]["runtime"][0] < 1:
        ax.yaxis.set_major_formatter(("{x:.3f}"))
    ax.tick_params(axis="x", labelrotation=90)

    if system_name == "juwelsbooster":
        system_name_string = "JUWELS Booster"
    elif system_name == "jupiter":
        system_name_string = "JUPITER"
    elif system_name == "juredadc":
        system_name_string = "JURECA DC"
    else:
        raise ValueError("Unknown system name: " + system_name)


    ax.set_xlabel("Number of " +  system_name_string + " Nodes (log-scale)")
    ax.set_ylabel("Runtime / s (log-scale)")
    plot.legend(loc="lower right", bbox_to_anchor=(legend_position[0], legend_position[1]))
    plot.savefig(os.path.join(output_dir, "plot.png"), bbox_inches="tight")
    plot.savefig(os.path.join(output_dir, "plot.pdf"), bbox_inches="tight")


def generate(system_name, input_dir, output_dir, skip_weak_scaling, legend_position):
    prepare_output_dir(input_dir, output_dir)
    plot_data = extract_data(input_dir)
    generate_plot_file(system_name, output_dir, plot_data, skip_weak_scaling, legend_position)
