# mypy: allow-untyped-defs
import argparse
import time

import numpy as np  # type: ignore[import]
import pandas as pd  # type: ignore[import]
from dlrm_s_pytorch import unpack_batch  # type: ignore[import]
from dlrm_utils import (  # type: ignore[import]
    dlrm_wrap,
    fetch_model,
    make_test_data_loader,
)

import torch


def run_forward(model, **batch):
    """The purpose of this function is to time the forward run of the model.
    The model forward happens a 100 times and each pass is timed. The average
    of this 100 runs is returned as avg_time.
    """
    time_list = []
    X, lS_o, lS_i = batch["X"], batch["lS_o"], batch["lS_i"]
    for _ in range(100):
        start = time.time()
        with torch.no_grad():
            model(X, lS_o, lS_i)
        end = time.time()
        time_taken = end - start
        # pyrefly: ignore  # bad-argument-type
        time_list.append(time_taken)
    avg_time = np.mean(time_list[1:])
    return avg_time


def make_sample_test_batch(raw_data_path, processed_data_path, device):
    """Create the test_data_loader and sample a batch from it. This batch will be used
    to measure the forward pass of the model throughout this experiment.
    """
    test_data_loader = make_test_data_loader(raw_data_path, processed_data_path)

    test_iter = iter(test_data_loader)

    test_batch = next(test_iter)

    X_test, lS_o_test, lS_i_test, _, _, _ = unpack_batch(test_batch)

    X, lS_o, lS_i = dlrm_wrap(X_test, lS_o_test, lS_i_test, device)
    batch = {"X": X, "lS_o": lS_o, "lS_i": lS_i}

    return batch


def measure_forward_pass(sparse_model_metadata, device, sparse_dlrm, **batch):
    """Measures and tracks the forward pass of the model for all the sparsity levels, block shapes and norms
    available in sparse_model_metadata file.
    If sparse_dlrm=True, then the SparseDLRM model is loaded, otherwise the standard one is.
    """
    time_taken_dict: dict[str, list] = {
        "norm": [],
        "sparse_block_shape": [],
        "sparsity_level": [],
        "time_taken": [],
    }

    metadata = pd.read_csv(sparse_model_metadata)

    for _, row in metadata.iterrows():
        norm, sbs, sl = row["norm"], row["sparse_block_shape"], row["sparsity_level"]
        model_path = row["path"]
        model = fetch_model(model_path, device, sparse_dlrm=sparse_dlrm)
        time_taken = run_forward(model, **batch)
        out_str = f"{norm}_{sbs}_{sl}={time_taken}"
        print(out_str)
        time_taken_dict["norm"].append(norm)
        time_taken_dict["sparse_block_shape"].append(sbs)
        time_taken_dict["sparsity_level"].append(sl)
        time_taken_dict["time_taken"].append(time_taken)

    time_df = pd.DataFrame(time_taken_dict)

    if sparse_dlrm:
        time_df["dlrm_type"] = "with_torch_sparse"
    else:
        time_df["dlrm_type"] = "without_torch_sparse"

    return time_df


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--raw-data-file", "--raw_data_file", type=str)
    parser.add_argument("--processed-data-file", "--processed_data_file", type=str)
    parser.add_argument("--sparse-model-metadata", "--sparse_model_metadata", type=str)

    args = parser.parse_args()

    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    print(device)

    batch = make_sample_test_batch(args.raw_data_file, args.processed_data_file, device)

    print("Forward Time for Sparse DLRM")
    sparse_dlrm_time_df = measure_forward_pass(
        args.sparse_model_metadata, device, sparse_dlrm=True, **batch
    )
    print(sparse_dlrm_time_df)

    print("Forward Time for Normal DLRM")
    norm_dlrm_time_df = measure_forward_pass(
        args.sparse_model_metadata, device, sparse_dlrm=False, **batch
    )
    print(norm_dlrm_time_df)

    forward_time_all = pd.concat([sparse_dlrm_time_df, norm_dlrm_time_df])
    forward_time_all.to_csv("dlrm_forward_time_info.csv", index=False)
