from matplotlib.ticker import FuncFormatter
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.lines import Line2D
from enum import Enum


class columns ():
    benchmark_type = 'benchmark_type'
    proc_num = 'proc_num'
    msg_size_bytes = 'msg_size_bytes'
    repetitions = 'repetitions'
    t_min_usec = 't_min_usec'
    t_max_usec = 't_max_usec'
    t_avg_usec = 't_avg_usec'
    mpi_datatype = 'mpi_datatype'
    mpi_red_datatype = 'mpi_red_datatype'
    mpi_red_op = 'mpi_red_op'
    creation_time = 'creation_time'
    n_nodes = 'n_nodes'
    off_cache_flag = 'off_cache_flag'
    algorithm = 'algorithm'


class collectives(Enum):
    Bcast = 'Bcast'
    Reduce = 'Reduce'
    Allreduce = 'Allreduce'
    Alltoall = 'Alltoall'
    Scatter = 'Scatter'
    Reduce_scatter = 'Reduce_scatter'
    Allgather = 'Allgather'
    Gather = 'Gather'


data_file = "./data/data_04_11_25_algs.csv"

df_multinode = pd.read_csv(data_file, delimiter=',')
df_multinode.fillna(0, inplace=True)
df_multinode = df_multinode[df_multinode[columns.off_cache_flag] == 50]
for c in collectives:

    df_single = df_multinode[df_multinode[columns.benchmark_type]
                             == c.value]
    df_single = df_single[df_single[columns.msg_size_bytes] > 1000]
    df_single = df_single[[columns.proc_num, columns.msg_size_bytes,
                           columns.t_avg_usec, columns.algorithm]]

    df_gather_best = df_single.loc[  # pyright: ignore
        df_single.groupby([columns.msg_size_bytes, columns.proc_num])[  # pyright: ignore
            columns.t_avg_usec].idxmin()
    ].reset_index(drop=True)

    df_gather_best = df_gather_best[df_gather_best[columns.msg_size_bytes] > 2**17]

    pivot_best = df_gather_best.pivot(index=columns.msg_size_bytes,  # pyright: ignore
                                      columns=columns.proc_num, values=columns.t_avg_usec)

    X = pivot_best.columns.values   # proc_num
    Y = pivot_best.index.values     # msg_size_bytes
    X, Y = np.meshgrid(X, Y)   # pyright: ignore
    Z = pivot_best.values

    alg_pivot = df_gather_best.pivot(
        index=columns.msg_size_bytes,
        columns=columns.proc_num,
        values=columns.algorithm
    )

    algorithms = alg_pivot.values.flatten()
    unique_algs = sorted(pd.unique(algorithms[~pd.isna(algorithms)]))
    color_map = {alg: i for i, alg in enumerate(unique_algs)}
    color_values = np.array([color_map.get(a, np.nan) for a in algorithms])

    fig = plt.figure(figsize=(16, 9))
    ax = fig.add_subplot(111, projection='3d')
    surf = ax.plot_wireframe(X, Y, Z,
                             color='black', linewidths=1)
    surf_points = ax.scatter(X,
                             Y, Z, c=color_values, cmap='viridis', s=20, depthshade=False)  # pyright: ignore

    handles = [
        Line2D([0], [0],
               marker='o', color='w',
               label=alg,
               markerfacecolor=plt.cm.viridis(
                   color_map[alg] / max(len(unique_algs)-1, 1)),
               markersize=8)
        for alg in unique_algs
    ]

    ax.legend(handles=handles, title="Algorithm", loc='upper right')
    ax.set_xlabel("Process Count")
    ax.set_ylabel("Message Size [B]")
    ax.set_zlabel("Average Time [μs]")
    ax.set_title(f"{c.value}")
    ax.set_xticks(pivot_best.columns.values)  # pyright: ignore
    ax.set_xticklabels(pivot_best.columns.values)
    ax.set_yticks(Y[:, 0])
    ymin, ymax = ax.get_ylim()
    ax.set_ylim(ymin*0.8, ymax)   # 30% more space at top
    ax.set_yticklabels([f"$2^{{{int(np.log2(v))}}}$" for v in Y[:, 0]])
    plt.savefig(f"./plots/{c.value.lower()}_algcomp.png")