#!/usr/bin/env python3
"""
This file parses the JSON output of Rocprof to get aggregated statistics for
GPU performance analysis.
Don't forget to load the cray-python module before running the script!
"""

import json
import numpy as np
import pandas as pd
from texttable import Texttable
import latextable as lt
from pathlib import Path
import seaborn as sns
from matplotlib import pyplot as plt
import argparse
import sys


def _parse_args() -> argparse.Namespace:
    """Parse user arguments for generating latex tables from CSV results."""
    # Setup argparser.
    parser = argparse.ArgumentParser(
                    prog="create_tables_latex.py",
                    description="Creates LaTeX tables from csv files from " +
                                "experiments.",
                    )
    parser.add_argument("-t", '--table', action="store_true", 
                        help="set to generate latex table.")
    parser.add_argument("-f", '--figure', action="store_true", 
                        help="set to generate boxplot.")
    parser.add_argument("-ps", '--print_stats', action="store_true", 
                        help="set to print statistics.")
    parser.add_argument("-ts", '--table_stats', action="store_true", 
                        help="set to generate statistics table.")
    parser.add_argument("-s", '--speedup', action="store_true", 
                        help="set to generate speedup table.")
    
    # Parse args and just print usage in case no arguments are given.
    args = parser.parse_args(args=None if sys.argv[1:] else ['--help'])

    return args


def _load_json(folder, prefix, iteration, nlat, nlon):
    """
    Load JSON output and slice into the requested iteration.
    Iterations are 1-based numbers. 0 represents the data of the initialization.
    """
    start_kernel = "( kernel(eltinv$eltinv_mod_$ck_L136_2_cce$noloop$form)"
    
    # Read JSON file.
    with open(f"../profiling/{folder}/{prefix}_nlat_{nlat}_nlon_{nlon}/" +
               "rocprof.json") as fp:
        data = json.load(fp)
        data = np.array(data["traceEvents"])
                
    # Variables for tracking the iteration, and storing start and end indexes 
    # of the requested itteration.
    cur_itter = 0
    start_index = -1
    end_index = -1
    
    # Loop over the JSON data until requested iteration is found.
    for i, dict_item in enumerate(iter(data)):
        if not dict_item.keys() \
           or "args" not in dict_item \
           or "args" not in dict_item["args"]:
            continue
        
        if dict_item["name"].startswith("hipModuleLaunchKernel") \
           and dict_item["args"]["args"].startswith(start_kernel):
            cur_itter += 1
            
            # Return initialization meta data if requested iteration is 0.
            if iteration == 0:
                start_index = 0
                end_index = i
                break
            # Otherwise create
            elif cur_itter == iteration:
                start_index = i
            elif cur_itter == iteration + 1:
                end_index = i
                break

    # Return slice of the JSON file which only contains requested itteration.
    # Return entire set in case the start_kernel is missing.
    if start_index == -1:
        return data

    return data[start_index:end_index]


def constr_memory_df(data):
    """Construct a DataFrame with memory operations within the provided data."""
    # First construct a list of dicts, than convert into DataFrame later on.
    df_list = []
    
    # Loop over all operations.
    for dict_item in iter(data):
        if not dict_item.keys() \
           or "args" not in dict_item \
           or "args" not in dict_item["args"]:
            continue
        
        # Store HtoD and DtoH memory copies.
        if dict_item["name"] in ["hipMemcpyHtoD", "hipMemcpyDtoH"]:
            df_list.append({
                "Metric": dict_item["name"],
                # For size, get sizeBytes section, then slice off constant str.
                "SizeBytes": dict_item["args"]["args"].split(" ")[3][10:-2],
                "DurationNs": dict_item["args"]["DurationNs"],
            })
            
    # Convert list of dicts into a DataFrame, set types, and return.
    df = pd.DataFrame(df_list)
    df[["SizeBytes", "DurationNs"]] = df[["SizeBytes", "DurationNs"]].apply(pd.to_numeric)
    return df


def constr_memory_table(nlat_nlons, mem_dfs, iteration):
    """
    Construct and return a LaTeX table with min, mean, max, sum statistics of
    memory operations in the provided iteration.
    """
    # Setup table with a column per nlat_nlon configuration and row per metric.
    table = Texttable()
    align = ["c" for _ in nlat_nlons]
    align.insert(0, "l")
    table.set_cols_align(align)
    table.set_cols_valign(["m" for _ in align])
    table.set_precision(2)
    
    # Setup table header and fill in table.
    header = [f"({nlat},{nlon})" for (nlat, nlon) in nlat_nlons]
    header.insert(0, "Metric")
    rows = [header]
    
    opperations = [
        pd.Series.min,
        pd.Series.mean,
        pd.Series.max,
        pd.Series.sum,
    ]
    
    # Add rows for the HostToDevice transfers.
    header_row = [0] * len(nlat_nlons)
    header_row.insert(0, "Transfer size HostToDevice")
    rows.append(header_row)
    
    for opp in opperations:
        new_row = [opp(mem_df[mem_df.Metric == "hipMemcpyHtoD"]["SizeBytes"]) 
                   for mem_df in mem_dfs]
        new_row.insert(0, opp.__name__)
        rows.append(new_row)

    # Add rows for the DeviceToHost transfers.
    header_row = [0] * len(nlat_nlons)
    header_row.insert(0, "Transfer size DeviceToHost")
    rows.append(header_row)
    
    for opp in opperations:
        new_row = [opp(mem_df[mem_df.Metric == "hipMemcpyDtoH"]["SizeBytes"]) 
                   for mem_df in mem_dfs]
        new_row.insert(0, opp.__name__)
        rows.append(new_row)

    # Transform and return table in Latex format.
    table.add_rows(rows)
    caption = "Minimum, mean, maximum, and sum of transfer times for " + \
              f"HostToDevice and DeviceToHost of iteration {iteration}. " + \
              " All numbers are in MB."
    return lt.draw_latex(table, 
                         caption=caption, 
                         label=f"tab:ectrans_memory_size_stats",
                         caption_above=False,
                         position="H")


def plot_memory_violin(nlat_nlons, mem_dfs):
    """
    Create a violin plot of the bytes send for each workload for the 
    provided iteration.
    """
    # Add nlat_nlon config columns to the dfs and then merge them into one.
    for i in range(len(mem_dfs)):
        mem_dfs[i]["nlat_nlon"] = [f"({nlat_nlons[i][0]},{nlat_nlons[i][1]})"] \
                                    * len(mem_dfs[i])
    df = pd.concat(mem_dfs, ignore_index=True)
    
    # Setup figure and plot in violin style.
    ax = plt.figure(figsize=(9,7)).gca()
    sns.boxplot(data=df, x="nlat_nlon", y="SizeBytes", hue="Metric")
    ax.set_yscale('log')
    ax.set_xlabel("Workload size")
    ax.set_ylabel("Bytes Transferred")
    plt.savefig("figures/workload_tests/boxplot_transfers.png")
    plt.show()


def print_memory_stats(nlat_nlons, mem_dfs):
    """Print notable statistics on the memory transfers."""
    # Print first for H2D and then for D2H
    for metric in ["hipMemcpyHtoD", "hipMemcpyDtoH"]:
        print(f"\n\n\nVALUES FOR METRIC {metric}")
        for nlat_nlon, mem_df in zip(nlat_nlons, mem_dfs):
            unique_cnts = mem_df[mem_df["Metric"] == metric]["SizeBytes"].value_counts()
            small_unique_cnts = unique_cnts[unique_cnts.index < 100000]
            
            print("\nNLAT_NLON:  ", nlat_nlon)
            print("General stats:")
            print(f"\tValues: [{unique_cnts.index.min()}, " +
                f"{unique_cnts.index.max()}] with in total " +
                f"{unique_cnts.sum()} items.")
            print("Stats on smaller than 100KB:")
            print(f"\tValues: [{small_unique_cnts.index.min()}, " +
                f"{small_unique_cnts.index.max()}] with in total " +
                f"{small_unique_cnts.sum()} items.")
        

def constr_memory_stats_table(nlat_nlons, mem_dfs, iteration, 
                              metric="hipMemcpyHtoD"):
    """Create LaTeX table with statistics on small H2D copies."""
    # Setup table with a column per nlat_nlon configuration and row per metric.
    table = Texttable()
    align = ["c" for _ in nlat_nlons]
    align.insert(0, "l")
    table.set_cols_align(align)
    table.set_cols_valign(["m" for _ in align])
    table.set_precision(2)
    
    # Setup table header.
    header = [f"({nlat},{nlon})" for (nlat, nlon) in nlat_nlons]
    header.insert(0, "Metric")
    
    # Construct the rows to add.
    min_row = ["Min. (MB)"]
    max_row = ["Max. (MB)"]    
    sum_row = ["Count"]
    
    for mem_df in mem_dfs:
        unique_cnts = mem_df[mem_df["Metric"] == metric]["SizeBytes"].value_counts()
        small_unique_cnts = unique_cnts[unique_cnts.index < 100000]
        min_row.append(small_unique_cnts.index.min())
        max_row.append(small_unique_cnts.index.max())
        sum_row.append(small_unique_cnts.sum())

    # Fill table, transform, and return table in Latex format.
    table.add_rows([header, min_row, max_row, sum_row])
    caption = "Statistics on the HostToDevice memory copies smaller than " + \
              f"100KB for iteration {iteration}. "
    return lt.draw_latex(table, 
                         caption=caption, 
                         label=f"tab:ectrans_small_h2d_stats",
                         caption_above=False,
                         position="H")


def _get_iteration_duration(iteration, data):
    """Get the duration of the given iteration by finding the RocTX marker."""
    # Loop over all operations.
    for dict_item in iter(data):
        if not dict_item.keys() \
           or "name" not in dict_item:
            continue

        # Store duration of first item with the iteration marker name.
        if dict_item["name"] == "spectral_transform_7":
            return dict_item["dur"]

    # Return None in case marker was not found.
    return None


def constr_speedup_df(nlat_nlons, iteration, gpu_data, cpu_data):
    """
    Construct a DataFrame with execution times and speedup within the provided 
    data.
    """
    # First construct a list of dicts, than convert into DataFrame later on.
    # Add nlat_nlon column first.
    df_list = [{"nlat_nlon": nlat_nlon} for nlat_nlon in nlat_nlons]
    
    # For each nlat_nlon, add the time step duration into the column.
    for i, (gpu_d, cpu_d) in enumerate(zip(gpu_data, cpu_data)):
        df_list[i]["gpu_time"] = _get_iteration_duration(iteration, gpu_d)
        df_list[i]["cpu_time"] = _get_iteration_duration(iteration, cpu_d)
    
    # Convert list of dicts into a DataFrame, set types, and add speedup column.
    df = pd.DataFrame(df_list)
    df[["gpu_time", "cpu_time"]] = df[["gpu_time", "cpu_time"]] \
        .apply(pd.to_numeric)
    df["speedup"] = df["cpu_time"].div(df["gpu_time"])
        
    return df


def constr_speedup_table(nlat_nlons, speedup_df, iteration):
    """Create LaTeX table with speedup numbers of a given iteration."""
    # Setup table with a column per nlat_nlon configuration and row per metric.
    table = Texttable()
    header =["Workload size", "GPU runtime", "CPU runtime", "Speedup"]
    align = ["c" for _ in header[1:]]
    align.insert(0, "l")
    table.set_cols_align(align)
    table.set_cols_valign(["m" for _ in align])
    table.set_precision(2)

    # Construct rows to add.
    rows = [header]
    for _, row in speedup_df.iterrows():
        rows.append(list(row))

    # Fill table, transform, and return table in Latex format.
    table.add_rows(rows)
    caption = f"Speedup of porting ecTrans to GPU for iteration {iteration}."
    return lt.draw_latex(table,
                         caption=caption,
                         label=f"tab:ectrans_speedup_iter_7",
                         caption_above=False,
                         position="H")


if __name__ == "__main__":
    # Parse user arguments.
    args = _parse_args()
    
    # Load and slice the JSON file.
    nlat_nlons = [(128, 128), (256, 256), (512, 512), (1024, 1024), 
                  (1024, 1536), (1536, 1024), (1280, 1280)]
    iteration = 7
    
    if args.table or args.figure or args.print_stats or args.table_stats:
        memory_data = [_load_json("workload_tests", "res-gpu_prof", iteration, 
                                *nlat_nlon) 
                for nlat_nlon in nlat_nlons]

        # Construct DataFrame with memory operations. 
        mem_dfs = [constr_memory_df(data) for data in memory_data]
    
    # Create transfer size table and store.
    if args.table:
        latex = constr_memory_table(nlat_nlons, mem_dfs, iteration)
        Path("tables/latex/workload_tests").mkdir(parents=True, 
                                                  exist_ok=True)
        with open(Path("tables/latex/workload_tests/transfer_size_stats.tex"), 
                "w") as fp:
            fp.write(latex + "\n")
        print(latex, end="\n\n")

    # Plot violin plot of the memory transfers.
    if args.figure:
        plot_memory_violin(nlat_nlons, mem_dfs)

    # Print statistics on the memory transfers.
    if args.print_stats:
        print_memory_stats(nlat_nlons, mem_dfs)
        
    # Print statistics on the memory transfers.
    if args.table_stats:
        # Create table for HostToDevice transfers.
        latex = constr_memory_stats_table(nlat_nlons, mem_dfs, iteration)
        Path("tables/latex/workload_tests").mkdir(parents=True, 
                                                  exist_ok=True)
        with open(Path("tables/latex/workload_tests/small_h2d_stats.tex"), 
                  "w") as fp:
            fp.write(latex + "\n")
        print(latex, end="\n\n")

        # Create table for DeviceToHost transfers
        latex = constr_memory_stats_table(nlat_nlons, mem_dfs, iteration,
                                          metric="hipMemcpyDtoH")
        Path("tables/latex/workload_tests").mkdir(parents=True, 
                                                  exist_ok=True)
        with open(Path("tables/latex/workload_tests/small_d2h_stats.tex"), 
                  "w") as fp:
            fp.write(latex + "\n")
        print(latex, end="\n\n")
        
    # Create speedup table.
    if args.speedup:
        # Only fetch initialization data by requesting iteration 0.
        gpu_data = [_load_json("workload_tests", "res-gpu_prof", 0,
                               *nlat_nlon) 
                    for nlat_nlon in nlat_nlons]
        cpu_data = [_load_json("speedup_tests", "res-cpu_prof", 0,
                               *nlat_nlon) 
                    for nlat_nlon in nlat_nlons]
        speedup_df = constr_speedup_df(nlat_nlons, iteration, gpu_data, cpu_data)
        latex = constr_speedup_table(nlat_nlons, speedup_df, iteration)
        Path("tables/latex/speedup_tests").mkdir(parents=True,
                                                  exist_ok=True)
        with open(Path("tables/latex/speedup_tests/speedup_dp.tex"),
                  "w") as fp:
            fp.write(latex + "\n")
        print(latex, end="\n\n")
