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

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.")
    
    # 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
    itter_df = df.iloc[start_idxs[itteration - 1]:start_idxs[itteration]]
    
    return itter_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.
    itter_dfs = [_get_basic_profile_data(nproma, iteration) 
                 for nproma in npromas]
    
    # Extract data of specific kernel for each nproma.
    itter_dfs = [df.loc[df["KernelName"] == kernelName] for df in itter_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", "L2CacheHit"]
    
    for metric in metrics:
        # Load values of current metric for each nproma value.
        new_row = [df.iloc[0][metric] for df in itter_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.
    print(rows)
    table.add_rows(rows)
    caption = f"Metrics for the {kernelName} 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
    itter_df = df.iloc[start_idxs[itteration - 1]:start_idxs[itteration]]
    
    return itter_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.
    itter_dfs = [_get_workload_profile_data(nlat, nlon, iteration) 
                 for (nlat, nlon) in nlat_nlons]
    
    # Extract data of specific kernel for each nproma.
    itter_dfs = [df.loc[df["KernelName"] == kernelName] for df in itter_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"({nlat},{nlon})" for (nlat, nlon) in nlat_nlons]
    header.insert(0, "Metric")
    rows = [header]
    
    # Add values for each metric selected.
    metrics = ["vgpr", "sgpr", "Wavefronts", "grd", "wgr", "VALUUtilization", 
               "VALUBusy", "MemUnitBusy", "L2CacheHit"]
    
    for metric in metrics:
        # Load values of current metric for each nproma value.
        new_row = [df.iloc[0][metric] for df in itter_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 = f"Metrics for the {kernelName} 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.
    itter_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.
        new_row = [df.iloc[0][metric] for df in itter_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")


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)]
    itteration = 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, itteration)
            Path("tables/latex/basic_tests").mkdir(parents=True, exist_ok=True)
            with open(Path(f"{dirs}/workload_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, 
                                                    itteration)
            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")
