#!/usr/bin/env python3
"""
This file generates word tables from the profiling data in the results folders.
It uses the docx package to create new word files with these tables and pandas
for the data mangement:
 - pip install --user python-docx pandas
Don't forget to load the cray-python module before running the script!
"""

import docx
import pandas as pd


def get_profile_data(nproma: int, itteration: int) -> pd.DataFrame:
    """Load the profiling data from the results and aggregate hardware counters."""
    df = pd.read_csv(f"../profiling/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.iloc[start_idxs[itteration - 1]:start_idxs[itteration]]
    
    return itter


def add_metrics_table(doc: docx.Document, kernelName: str):
    """Adds new section with table containing transfer data."""
    # Load the data from the 7th iteration for each nproma value.
    df_nproma8 = get_profile_data(8, 7)
    # df_nproma16 = get_profile_data(16, 7)
    # df_nproma32 = get_profile_data(32, 7)
    # df_nproma64 = get_profile_data(64, 7)
    # df_nproma128 = get_profile_data(128, 7)
    # df_nproma256 = get_profile_data(256, 7)
    
    # Extract the kernel to create the table for.
    df_nproma8 = df_nproma8.loc[df_nproma8["KernelName"] == kernelName]
    # df_nproma16 = df_nproma16.loc[df_nproma16["KernelName"] == kernelName]
    # df_nproma32 = df_nproma32.loc[df_nproma32["KernelName"] == kernelName]
    # df_nproma64 = df_nproma64.loc[df_nproma64["KernelName"] == kernelName]
    # df_nproma128 = df_nproma128.loc[df_nproma128["KernelName"] == kernelName]
    # df_nproma256 = df_nproma256.loc[df_nproma256["KernelName"] == kernelName]
    
    print(df_nproma8)
    
    # Add header to document for new table.
    doc.add_heading(f"Metrics table - {kernelName}", 0)
    
    # Add table and fill in header.
    table = doc.add_table(rows=1, cols=7)
    header = ["Metric", "nproma=8", "nproma=16", "nproma=32", "nproma=64", 
              "nproma=128", "nproma=256"]
    t_head = table.rows[0].cells
    
    for row, value in zip(t_head, header):
        row.text = value

    # Add values for each metric selected.
    metrics = ["vgpr", "sgpr", "Wavefronts", "grd", "wgr", "VALUUtilization", 
               "VALUBusy", "MemUnitBusy", "L2CacheHit"]
    
    for metric in metrics:
        row = table.add_row().cells
        row[0] = metric
        row[1] = df_nproma8
    

if __name__ == "__main__":
    # Create new document.
    doc = docx.Document()

    # Add metrics table.
    add_metrics_table(doc, "easre1b$easre1b_mod_$ck_L75_1")
    add_metrics_table(doc, "efourier_in$efourier_in_mod_$ck_L61_2_cce$noloop$form")
    add_metrics_table(doc, "trltog_cudaaware$trltog_mod_$ck_L559_4")
    add_metrics_table(doc, "trgtol_cudaaware$trgtol_mod_$ck_L563_3")
    add_metrics_table(doc, "efourier_out$efourier_out_mod_$ck_L56_1")
    add_metrics_table(doc, "eprfi2b$eprfi2b_mod_$ck_L84_1_cce$noloop$form")    
    
    # Save the document in the tables folder.
    doc.save("tables/ecTrans_dwarf_tables.docx")
