#!/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


def _load_json(iteration, nlat, nlon):
    """Load JSON output and slice into the requested iteration."""
    start_kernel = "( kernel(eltinv$eltinv_mod_$ck_L136_2_cce$noloop$form)"
    
    # Read JSON file.
    with open(f"../profiling/workload_tests/res-gpu_prof_nlat_{nlat}_" + 
              f"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 itteration 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
            
            if 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 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 requested 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])
    
    # 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 = [mem_df[mem_df.Metric == "hipMemcpyHtoD"]["SizeBytes"].opp() 
                   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 = [mem_df[mem_df.Metric == "hipMemcpyDtoH"]["SizeBytes"].opp() 
                   for mem_df in mem_dfs]
        new_row.insert(0, opp.__name__)
        rows.append(new_row)

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


if __name__ == "__main__":
    # Load and slice the JSON file.
    nlat_nlons = [(128, 128), (256, 256), (512, 512), (1024, 1024), 
                  (1024, 1536), (1536, 1024), (1280, 1280)]
    iteration = 7    
    data = [_load_json(iteration, *nlat_nlon) for nlat_nlon in nlat_nlons]
    
    # Construct DataFrame with memory operations. 
    mem_dfs = [constr_memory_df(d) for d in data]
    
    # Create transfer size table and store.
    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")
