#!/usr/bin/env python3
"""
This file generates word 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
from texttable import Texttable
from typing import List, Tuple
from pathlib import Path


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):
    """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])
    
    # 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:{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}_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):
    """Adds new section with table containing transfer data."""
    # Load the data from the given iteration for each nproma 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 npromas]
    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"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:{kernelName}",
                         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]
    itteration = 7

    # 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:
        # 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")

        # Generate and store metrics table for each kernel from the 
        # workload_test folder.
        latex = generate_workload_metrics_table(kernel, itteration)
        Path("tables/latex/workload_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")
