#!/usr/bin/env python3

import argparse
import glob
import subprocess
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd

pd.set_option("display.float_format", lambda x: f"{x:.2f}")

AIC_KEYS = [
    "aic_time(us)",
    "aic_cube_time(us)",
    "aic_cube_ratio",
    "aic_mte2_time(us)",
    "aic_mte2_ratio",
    "aic_fixpipe_time(us)",
    "aic_fixpipe_ratio",
    "aic_fixpipe_active_bw(GB/s)",
    "aic_main_mem_read_bw(GB/s)",
    "aic_main_mem_write_bw(GB/s)",
    "aic_read_hit_rate(%)",
    "aic_write_hit_rate(%)",
    "aic_total_hit_rate(%)",
]

AIV_KEYS = [
    "aiv_time(us)",
    "aiv_vec_time(us)",
    "aiv_mte2_time(us)",
    "aiv_mte2_ratio",
    "aiv_mte2_active_bw(GB/s)",
    "aiv_mte3_time(us)",
    "aiv_mte3_ratio",
    "aiv_mte3_active_bw(GB/s)",
    "aiv_gm_to_ub_bw(GB/s)",
    "aiv_ub_to_gm_bw(GB/s)",
    "aiv_main_mem_read_bw(GB/s)",
    "aiv_main_mem_write_bw(GB/s)",
    "aiv_read_hit_rate(%)",
    "aiv_write_hit_rate(%)",
    "aiv_total_hit_rate(%)",
]

def get_aic_aiv_info(path) -> tuple[pd.Series, pd.Series]:
    df = pd.read_csv(path).fillna(0)
    aic_df = df[df["sub_block_id"].str.startswith("cube")].loc[:, df.columns.str.startswith("aic")]
    aic_info = pd.DataFrame(index=aic_df.columns, columns=["min", "mean", "max", "(max-min)/mean"])
    aic_info["min"] = aic_df.min()
    aic_info["mean"] = aic_df.mean()
    aic_info["max"] = aic_df.max()
    aic_info["(max-min)/mean"] = (aic_info["max"] - aic_info["min"]) / aic_info["mean"]

    aiv_df = df[df["sub_block_id"].str.startswith("vector")].loc[:, df.columns.str.startswith("aiv")]
    aiv_info = pd.DataFrame(index=aiv_df.columns, columns=["min", "mean", "max", "(max-min)/mean"])
    aiv_info["min"] = aiv_df.min()
    aiv_info["mean"] = aiv_df.mean()
    aiv_info["max"] = aiv_df.max()
    aiv_info["(max-min)/mean"] = (aiv_info["max"] - aiv_info["min"]) / aiv_info["mean"]
    return aic_info, aiv_info


def get_user():
    ret = subprocess.run("whoami", stdout=subprocess.PIPE, text=True, shell=True)
    return ret.stdout.strip()


def run_msprof(args: list[str]):
    output_path = "./msprof_output/" + datetime.now().strftime("%Y%m%d%H%M%S")
    output_path = Path(output_path)
    all_args = [
        "msprof",
        "op",
        "--launch-count=10",
        f"--output={output_path}",
        *args,
    ]
    cmd_str = " ".join(all_args)
    result = subprocess.run(cmd_str, shell=True)
    opinfo_path_list = glob.glob(f"{output_path}/OPPROF*/OpBasicInfo*.csv")
    if len(opinfo_path_list) == 0:
        print("no OpBasicInfo csv found")
    elif len(opinfo_path_list) > 1:
        print("more than 1 OpBasicInfo.csv, pass")
    else:
        print()
        opbasic_info_path = Path(opinfo_path_list[0])

        opbasic_info_df = pd.read_csv(opbasic_info_path)

        pipe_path = glob.glob(f"{opbasic_info_path.parent}/PipeUtilization*.csv")[0]
        aic_pipe_info, aiv_pipe_info = get_aic_aiv_info(pipe_path)
        memory_path = glob.glob(f"{opbasic_info_path.parent}/Memory*.csv")[0]
        aic_memory_info, aiv_memory_info = get_aic_aiv_info(memory_path)
        l2cache_path = glob.glob(f"{opbasic_info_path.parent}/L2Cache*.csv")[0]
        aic_l2cache_info, aiv_l2cache_info = get_aic_aiv_info(l2cache_path)

        aic_info = pd.concat([aic_pipe_info, aic_memory_info, aic_l2cache_info]).drop_duplicates()
        aiv_info = pd.concat([aiv_pipe_info, aiv_memory_info, aiv_l2cache_info]).drop_duplicates()
        aic_info = aic_info[aic_info.index.isin(AIC_KEYS)].loc[AIC_KEYS]
        aiv_info = aiv_info[aiv_info.index.isin(AIV_KEYS)].loc[AIV_KEYS]
        print(aic_info, "\n")
        print(aiv_info, "\n")
        print(opbasic_info_df, "\n")
        print(pipe_path)
        print(memory_path)

    if get_user() == "root":
        subprocess.run(f"chmod -R 777 {output_path}", shell=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("args", nargs="*", type=str)
    args = parser.parse_args()
    run_msprof(args.args)