#!/usr/bin/env python3
"""
This file generates LaTeX tables from the profiling data in the results folders.
It uses the 'latextable' package to create the Latex syntax tables, and 'pandas'
for the data mangement:
 - pip install --user texttable latextable pandas
Don't forget to load the cray-python module before running the script!
"""

from pylatex.utils import escape_latex
import latextable as lt
import pandas as pd
import numpy as np
from texttable import Texttable
from typing import List, Tuple
from pathlib import Path
import sqlite3
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("-b", '--basic', action="store_true", 
                        help="set to generate basic test tables.")
    parser.add_argument("-w", '--workload', action="store_true", 
                        help="set to generate workload test tables.")
    parser.add_argument("-wm", '--workload_memory', action="store_true", 
                        help="set to generate workload test memory " +
                             "transfer tables.")
    parser.add_argument("-wl", '--workload_ldsbank', action="store_true",
                        help="set to generate workload test LDSBankConflict " +
                             "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 _get_basic_profile_data(nproma: int, itteration: int) -> pd.DataFrame:
    """
    Load the profiling data from the basic results and aggregate hardware 
    counters.
    """    
    df = pd.read_csv(f"../profiling/basic_tests/res-gpu_prof_nproma_{nproma}" +
                      "/rocprof.csv")

    # Get starting indexes of each iteration, and filter out rows of 7th 
    # iteration.
    start_kernel = "eltinv$eltinv_mod_$ck_L136_2_cce$noloop$form.kd"
    start_idxs = df.loc[df["KernelName"] == start_kernel].index
    iter_df = df.iloc[start_idxs[itteration - 1]:start_idxs[itteration]]
    
    return iter_df


def generate_basic_metrics_table(
        kernelName: str, 
        npromas: List[int], 
        iteration: int,
    ) -> str:
    """Adds new section with table containing transfer data."""
    # Load the data from the given iteration for each nproma value.
    iter_dfs = [_get_basic_profile_data(nproma, iteration)
                 for nproma in npromas]
    
    # Extract data of specific kernel for each nproma.
    iter_dfs = [df.loc[df["KernelName"] == kernelName] for df in iter_dfs]

    # Setup table with a column per nproma and row per metric.
    table = Texttable()
    align = ["c" for _ in npromas]
    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"nproma={nproma}" for nproma in npromas]
    header.insert(0, "Metric")
    rows = [header]
    
    # Add values for each metric selected.
    metrics = ["vgpr", "sgpr", "Wavefronts", "grd", "wgr", "VALUUtilization",
               "VALUBusy", "MemUnitBusy", "MemUnitStalled", "L2CacheHit",
               "LDSBankConflict"]
    
    for metric in metrics:
        # Load values of current metric for each nproma value.
        # Format the strings to have wanted LateX styling.
        new_row = [f"\\num{{{df.iloc[0][metric]:.2f}}}"
                        if type(df.iloc[0][metric]) == np.float64
                        else f"\\num{{{df.iloc[0][metric]}}}"
                        for df in iter_dfs]

        # Prepend row with the metric name and then add row to table.
        new_row.insert(0, metric)
        rows.append(new_row)

    # Format kernelName to be latex printable.
    kernelNameCaption = escape_latex(kernelName)

    # Transform and return table in Latex format.
    table.add_rows(rows)
    caption = f"Metrics for the {kernelNameCaption} kernel."
    return lt.draw_latex(table, 
                         caption=caption, 
                         label=f"tab:basic_{kernelName}",
                         caption_above=False,
                         position="H")


def _get_workload_profile_data(
        nlat: int, 
        nlon: int, 
        itteration: int,
    ) -> pd.DataFrame:
    """
    Load the profiling data from the workload results and aggregate hardware 
    counters.
    """
    df = pd.read_csv(f"../profiling/workload_tests/res-gpu_prof_nlat_{nlat}_" +
                     f"nlon_{nlon}/rocprof.csv")

    # Get starting indexes of each iteration, and filter out rows of 7th 
    # iteration.
    start_kernel = "eltinv$eltinv_mod_$ck_L136_2_cce$noloop$form.kd"
    start_idxs = df.loc[df["KernelName"] == start_kernel].index
    iter_df = df.iloc[start_idxs[itteration - 1]:start_idxs[itteration]]
    
    # Set datatypes correctly that are inferred wrongly.
    iter_df = iter_df.convert_dtypes()

    return iter_df


def generate_workload_metrics_table(
        kernelName: str, 
        nlat_nlons: List[Tuple[int]], 
        iteration: int,
    ) -> str:
    """Adds table containing computation data."""
    # Load the data from the given iteration for each nlat_nlon value.
    iter_dfs = [_get_workload_profile_data(nlat, nlon, iteration)
                 for (nlat, nlon) in nlat_nlons]
    
    # Extract data of specific kernel for each nproma.
    iter_dfs = [df.loc[df["KernelName"] == kernelName] for df in iter_dfs]

    # Setup table with a column per nproma 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"\\textbf{{({nlat},{nlon})}}" for (nlat, nlon) in nlat_nlons]
    header.insert(0, "\\textbf{Metric}")
    rows = [header]
    
    # Add values for each metric selected.
    metrics = ["vgpr", "sgpr", "Wavefronts", "grd", "wgr", "VALUUtilization", 
               "VALUBusy", "MemUnitBusy", "MemUnitStalled", "L2CacheHit",
               "LDSBankConflict"]
    
    for metric in metrics:
        # Load values of current metric for each nproma value.
        # Format the strings to have wanted LateX styling, which depends on the
        # datatype of the variable.
        new_row = [f"\\num{{{df.iloc[0][metric]:.2f}}}"
                   if type(df.iloc[0][metric]) == np.float64
                   else f"\\num{{{df.iloc[0][metric]}}}"
                   for df in iter_dfs]
        
        # Prepend row with the metric name and then add row to table.
        new_row.insert(0, metric)
        rows.append(new_row)

    # Format kernelName to be latex printable.
    kernelNameCaption = f"\\texttt{{{escape_latex(kernelName)}}}"

    # Transform and return table in Latex format.
    table.add_rows(rows)
    caption = f"Metrics for the {kernelNameCaption} kernel."
    return lt.draw_latex(table, 
                         caption=caption, 
                         label=f"tab:comp_{kernelName}",
                         caption_above=False,
                         position="H")
    

def _get_workload_transfer_data(nlat: int, nlon: int) -> pd.DataFrame:
    """Load the tranfser data from the workload results."""
    # Load total kernel execution time from CSV.
    df = pd.read_csv(f"../profiling/workload_tests/res-gpu_prof_nlat_{nlat}_" +
                     f"nlon_{nlon}/rocprof.stats.csv")
    total_time = df["TotalDurationNs"].sum()
    
    # Load the number of calls, duration (ns), and transfer sizes (MB) from db.
    conn = sqlite3.connect(f"../profiling/workload_tests/res-gpu_prof_nlat_" +
                           f"{nlat}_nlon_{nlon}/rocprof.db")
    c = conn.cursor()
    c.execute("""
        SELECT
        Name as event,
        COUNT(*) as call_count,
        ROUND(SUM(DurationNs),2) as total_duration_ns,
        ROUND(SUM(SUBSTR( args, INSTR(args,"Bytes(")+6, INSTR(args,"))") - (INSTR(args,"Bytes(")+6)))/(1024.0*1024.0),2) as total_MBytes
        FROM HIP WHERE Name IN ("hipMemcpyHtoD","hipMemcpyDtoH")
        GROUP BY Name;
    """)

    # Store results into a Dict, to transfer later into DataFrame.
    data = {
        "Kernel time (ns)": [total_time], 
        "Call count HtoD": [], 
        "Call count DtoH": [], 
        "Transfer time HtoD (ns)": [], 
        "Transfer time DtoH (ns)": [], 
        "Transfer size HtoD (MB)": [], 
        "Transfer size DtoH (MB)": []
    }
    
    # Add metrics from the database.
    for row in c.fetchall():
        if row[0] == "hipMemcpyHtoD":
            data["Call count HtoD"].append(row[1])
            data["Transfer time HtoD (ns)"].append(row[2])
            data["Transfer size HtoD (MB)"].append(row[3])
        elif row[0] == "hipMemcpyDtoH":
            data["Call count DtoH"].append(row[1])
            data["Transfer time DtoH (ns)"].append(row[2])
            data["Transfer size DtoH (MB)"].append(row[3])

    
    # Create DataFrame and take average over the 10 iterations.
    df = pd.DataFrame(data=data)
    df["Call count HtoD"] = df["Call count HtoD"].div(10)
    df["Call count DtoH"] = df["Call count DtoH"].div(10)
    df["Transfer time HtoD (ns)"] = df["Transfer time HtoD (ns)"].div(10)
    df["Transfer time DtoH (ns)"] = df["Transfer time DtoH (ns)"].div(10)

    return df
    
    
def generate_workload_transfers_table(nlat_nlons: List[Tuple[int]]) -> str:
    """Adds new section with table containing transfer data."""
    # Load the hipTransfer times and calls from the CSV.
    iter_dfs = [_get_workload_transfer_data(nlat, nlon)
                 for (nlat, nlon) in nlat_nlons]

    # Setup table with a column per nproma 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]
    
    # Add values for each metric selected.
    metrics = ["Kernel time (ns)", "Call count HtoD", "Call count DtoH", 
               "Transfer time HtoD (ns)", "Transfer time DtoH (ns)", 
               "Transfer size HtoD (MB)", "Transfer size DtoH (MB)"]
    
    for metric in metrics:
        # Load values of current metric for each nproma value.
        # Format the strings to have wanted LateX styling, which depends on the
        # datatype of the variable.
        new_row = [f"\\num{{{df.iloc[0][metric]:.2f}}}"
                   if type(df.iloc[0][metric]) == np.float64
                   else f"\\num{{{df.iloc[0][metric]}}}"
                   for df in iter_dfs]
        
        # Prepend row with the metric name and then add row to table.
        new_row.insert(0, metric)
        rows.append(new_row)

    # Transform and return table in Latex format.
    table.add_rows(rows)
    caption = "Averages of 10 iterations of the kernel execution time and " + \
              "memory metrics for the various workload sizes."
    return lt.draw_latex(table, 
                         caption=caption, 
                         label=f"tab:ectrans_memory",
                         caption_above=False,
                         position="H")


def generate_workload_ldsbankconflict_table(nlat_nlons: List[Tuple[int]],
                                            iteration: int) -> str:
    """
    Creates table containing all kernels with LDSBankConflicts for a given
    iteration.
    """
    # Load the data from the given iteration for each nlat_nlon value.
    iter_dfs = [_get_workload_profile_data(nlat, nlon, iteration)
                 for (nlat, nlon) in nlat_nlons]

    # Filter on all kernels with non-zero LDSBankConflict values for any
    # workload size.
    iter_dfs = [iter_df[iter_df.LDSBankConflict > 0.0]
                [["KernelName", "LDSBankConflict"]]
                for iter_df in iter_dfs]

    # Get all unique indexes throughout the different workload sizes.
    indexes = [list(iter_df.index) for iter_df in iter_dfs]
    indexes = sorted(list({i for wl in indexes for i in wl}))

    # Setup table with a column per workload size and row per kernel.
    table = Texttable()
    align = ["c" for _ in range(len(nlat_nlons) + 1)]
    align.insert(1, "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"\\textbf{{({nlat},{nlon})}}" for (nlat, nlon) in nlat_nlons]
    header.insert(0, "\\textbf{ID}")
    header.insert(1, "\\textbf{Kernel}")
    rows = [header]

    # Add row for each kernel.
    for idx in indexes:
        new_row = [idx]

        # Add ID and KernelName.
        for iter_df in iter_dfs:
            if idx in iter_df.index:
                new_row.append(iter_df.loc[idx]["KernelName"])
                break

        # Add LDSBankConflict number, default to 0 if no data available.
        for iter_df in iter_dfs:
            if idx in iter_df.index:
                new_row.append(iter_df.loc[idx]["LDSBankConflict"])
            else:
                new_row.append(0)

        # Add newly constructed row to table data.
        rows.append(new_row)

    # Transform and return table in Latex format.
    table.add_rows(rows)
    caption = "All kernels with their occurence ID that have a non-zero" + \
                "\\texttt{LDSBankConflict} value."
    return lt.draw_latex(table,
                         caption=caption,
                         label=f"tab:ectrans_ldsbankconflict",
                         caption_above=False,
                         position="H")


if __name__ == "__main__":
    kernels = [
        "easre1b$easre1b_mod_$ck_L75_1.kd",
        "efourier_in$efourier_in_mod_$ck_L61_2_cce$noloop$form.kd",
        "trltog_cudaaware$trltog_mod_$ck_L559_4.kd",
        "trgtol_cudaaware$trgtol_mod_$ck_L563_3.kd",
        "efourier_out$efourier_out_mod_$ck_L56_1.kd",
        "eprfi2b$eprfi2b_mod_$ck_L84_1_cce$noloop$form.kd",
    ]
    npromas = [32, 64, 128, 192, 256]
    nlat_nlons = [(128, 128), (256, 256), (512, 512), (1024, 1024), 
                  (1024, 1536), (1536, 1024), (1280, 1280)]
    iteration = 7
    
    # Parse input arguments.
    args = _parse_args()

    # Create directory structure if missing.
    dirs = "tables/latex"
    Path("tables/latex").mkdir(parents=True, exist_ok=True)
    
    # Create a table for each of the kernels in the list and save in file.
    for kernel in kernels:
        if args.basic:
            # Generate and store metrics table for each kernel and nproma 
            # combination from the basic_tests folder.
            latex = generate_basic_metrics_table(kernel, npromas, iteration)
            Path("tables/latex/basic_tests").mkdir(parents=True, exist_ok=True)
            with open(Path(f"{dirs}/basic_tests/{kernel}.tex"), "w") as fp:
                fp.write(latex + "\n")
            print(latex, end="\n\n")

        if args.workload:
            # Generate and store metrics table for each kernel from the 
            # workload_test folder.
            latex = generate_workload_metrics_table(kernel, nlat_nlons, 
                                                    iteration)
            Path("tables/latex/workload_tests").mkdir(parents=True, 
                                                      exist_ok=True)
            with open(Path(f"{dirs}/workload_tests/comp_{kernel}.tex"), 
                      "w") as fp:
                fp.write(latex + "\n")
            print(latex, end="\n\n")
    
    # Generate and store memory transfers table for each kernel from 
    # the workload_test folder.
    if args.workload_memory:
        latex = generate_workload_transfers_table(nlat_nlons)
        Path("tables/latex/workload_tests").mkdir(parents=True, 
                                                  exist_ok=True)
        with open(Path(f"{dirs}/workload_tests/transfers_{kernel}.tex"), 
                  "w") as fp:
            fp.write(latex + "\n")
        print(latex, end="\n\n")

    # Generate and store LDSBankConflict table for each kernel from the
    # workload_test folder.
    if args.workload_ldsbank:
        latex = generate_workload_ldsbankconflict_table(nlat_nlons, iteration)
        Path("tables/latex/workload_tests").mkdir(parents=True,
                                                  exist_ok=True)
        with open(Path(f"{dirs}/workload_tests/LDSBankConflict.tex"),
                  "w") as fp:
            fp.write(latex + "\n")
        print(latex, end="\n\n")
