import pytest
import subprocess
from triton.profiler.viewer import get_min_time_flops, get_min_time_bytes, read, format_frames, derive_metrics, filter_frames, parse, apply_diff_profile
from triton.profiler.hooks.launch import COMPUTE_METADATA_SCOPE_NAME
import numpy as np

file_path = __file__
triton_example_file = file_path.replace("test_viewer.py", "examples/triton.json")
cuda_example_file = file_path.replace("test_viewer.py", "examples/cuda.json")
hip_example_file = file_path.replace("test_viewer.py", "examples/hip.json")
frame_example_file = file_path.replace("test_viewer.py", "examples/frame.json")
leaf_example_file = file_path.replace("test_viewer.py", "examples/leaf_nodes.json")


def test_help():
    # Only check if the viewer can be invoked
    subprocess.check_call(["proton-viewer", "-h"], stdout=subprocess.DEVNULL)


def test_exclusive_metrics():
    gf, inclusive_metrics, exclusive_metrics, device_info = read(triton_example_file)
    metrics = ["cpu_time/ns"]
    metrics = derive_metrics(gf, metrics, inclusive_metrics, exclusive_metrics, device_info)
    gf = filter_frames(gf, None, None, None, metrics[0])
    sorted_df = gf.dataframe.sort_values(by=[metrics[0]], ascending=False)
    actual = sorted_df.iloc[0:1]["name"].values[0]
    assert actual == "scope"


def test_sort():
    gf, inclusive_metrics, exclusive_metrics, device_info = read(leaf_example_file)
    gf = format_frames(gf, None)
    metrics = ["time/s", "time/ms", "time/us", "time/ns"]
    metrics = derive_metrics(gf, metrics, inclusive_metrics, exclusive_metrics, device_info)
    gf = filter_frames(gf, None, None, None, metrics[0])
    sorted_df = gf.dataframe.sort_values(by=[metrics[0]], ascending=False)
    actual = sorted_df.iloc[0:5]["name"].values
    expected = ["ROOT", "kernel_1_1_1", "kernel_3_1_1", "kernel_3_2_2", "kernel_1_2_2"]
    assert len(actual) == len(expected)
    assert all(a == b for a, b in zip(actual, expected))


@pytest.mark.parametrize("option", ["full", "file_function_line", "function_line", "file_function"])
def test_format_frames(option):
    gf, _, _, _ = read(frame_example_file)
    gf = format_frames(gf, option)
    if option == "full":
        idx = gf.dataframe["name"] == "/home/user/projects/example.py/test.py:1@foo"
    elif option == "file_function_line":
        idx = gf.dataframe["name"] == "test.py:1@foo"
    elif option == "function_line":
        idx = gf.dataframe["name"] == "1@foo"
    elif option == "file_function":
        idx = gf.dataframe["name"] == "test.py@foo"
    assert idx.sum() == 1


@pytest.mark.parametrize("option", ["include", "exclude"])
def test_filter_frames(option):
    include = ""
    exclude = ""
    gf, _, _, _ = read(frame_example_file)
    if option == "include":
        include = ".*test0.*"
    elif option == "exclude":
        exclude = ".*test1.*"
    gf = filter_frames(gf, include=include, exclude=exclude)
    idx = gf.dataframe["name"] == "test1"
    assert idx.sum() == 0
    idx = gf.dataframe["name"] == "test0"
    assert idx.sum() == 1


def test_filter_metadata():
    gf, _, _, _ = read(triton_example_file)
    assert COMPUTE_METADATA_SCOPE_NAME not in gf.dataframe["name"].tolist()
    assert "cuda_kernel" not in gf.dataframe["name"].tolist()
    assert "scope" in gf.dataframe["name"].tolist()
    assert "triton_kernel" in gf.dataframe["name"].tolist()


def test_parse():
    gf, derived_metrics = parse(["time/s"], triton_example_file)
    for derived_metric in derived_metrics:
        assert derived_metric in gf.inc_metrics or derived_metric in gf.exc_metrics


def test_min_time_flops():
    gf, _, _, device_info = read(cuda_example_file)
    ret = get_min_time_flops(gf.dataframe, device_info)
    device0_idx = gf.dataframe["device_id"] == "0"
    device1_idx = gf.dataframe["device_id"] == "1"
    device2_idx = gf.dataframe["device_id"] == "2"
    # sm89
    np.testing.assert_allclose(ret[device0_idx].to_numpy(), [[0.000025]], atol=1e-5)
    # sm90
    np.testing.assert_allclose(ret[device1_idx].to_numpy(), [[0.00005]], atol=1e-5)
    # sm100
    np.testing.assert_allclose(ret[device2_idx].to_numpy(), [[0.000025]], atol=1e-5)
    gf, _, _, device_info = read(hip_example_file)
    ret = get_min_time_flops(gf.dataframe, device_info)
    device0_idx = gf.dataframe["device_id"] == "0"
    device1_idx = gf.dataframe["device_id"] == "1"
    device2_idx = gf.dataframe["device_id"] == "2"
    # CDNA2
    np.testing.assert_allclose(ret[device0_idx].to_numpy(), [[0.000055]], atol=1e-5)
    # CDNA3
    np.testing.assert_allclose(ret[device1_idx].to_numpy(), [[0.000038]], atol=1e-5)
    # CDNA4
    np.testing.assert_allclose(ret[device2_idx].to_numpy(), [[0.000217]], atol=1e-5)


def test_min_time_bytes():
    gf, _, _, device_info = read(cuda_example_file)
    ret = get_min_time_bytes(gf.dataframe, device_info)
    device0_idx = gf.dataframe["device_id"] == "0"
    device1_idx = gf.dataframe["device_id"] == "1"
    # sm89
    np.testing.assert_allclose(ret[device0_idx].to_numpy(), [[9.91969e-06]], atol=1e-6)
    # sm90
    np.testing.assert_allclose(ret[device1_idx].to_numpy(), [[2.48584e-05]], atol=1e-6)
    gf, _, _, device_info = read(hip_example_file)
    ret = get_min_time_bytes(gf.dataframe, device_info)
    device0_idx = gf.dataframe["device_id"] == "0"
    device1_idx = gf.dataframe["device_id"] == "1"
    device2_idx = gf.dataframe["device_id"] == "2"
    # CDNA2
    np.testing.assert_allclose(ret[device0_idx].to_numpy(), [[3.125e-06]], atol=1e-6)
    # CDNA3
    np.testing.assert_allclose(ret[device1_idx].to_numpy(), [[1.93378e-05]], atol=1e-6)
    # CDNA4
    np.testing.assert_allclose(ret[device2_idx].to_numpy(), [[0.000125]], atol=1e-6)


def test_percentage():
    pass


def derivation_metrics_test(metrics, expected_data, sample_file, rtol=1e-7, atol=1e-6):
    gf, inclusive_metrics, exclusive_metrics, device_info = read(sample_file)
    assert len(inclusive_metrics + exclusive_metrics) > 0, "No metrics found in the input file"
    derived_metrics = derive_metrics(gf, metrics, inclusive_metrics, exclusive_metrics, device_info)
    for derived_metric in derived_metrics:
        np.testing.assert_allclose(gf.dataframe[derived_metric].to_numpy(), expected_data[derived_metric], rtol=rtol,
                                   atol=atol)


def test_avg_time_derivation():
    derivation_metrics_test(
        metrics=["avg_time/s", "avg_time/ms", "avg_time/us", "avg_time/ns"], expected_data={
            "avg_time/s (inc)": [0.0000512, 0.0000205, 0.000205,
                                 0.000205], "avg_time/ms (inc)": [0.0512, 0.02048, 0.2048, 0.2048], "avg_time/us (inc)":
            [51.2, 20.48, 204.8, 204.8], "avg_time/ns (inc)": [51200.0, 20480.0, 204800.0, 204800.0]
        }, sample_file=cuda_example_file)


def test_util():
    derivation_metrics_test(metrics=["util"], expected_data={
        "util": [np.nan, 0.247044, 0.147830, 0.118451],
    }, sample_file=cuda_example_file)


def test_time_derivation():
    derivation_metrics_test(
        metrics=["time/s", "time/ms", "time/us", "time/ns"], expected_data={
            "time/s (inc)": [0.000614, 0.0002048, 0.0002048, 0.0002048],
            "time/ms (inc)": [0.6144, 0.2048, 0.2048, 0.2048],
            "time/us (inc)": [614.4, 204.8, 204.8, 204.8],
            "time/ns (inc)": [614400.0, 204800.0, 204800.0, 204800.0],
            "time/% (inc)": [100.0, 50.0, 50.0, 50.0],
        }, sample_file=cuda_example_file)


def test_bytes_derivation():
    derivation_metrics_test(
        metrics=["byte/s", "gbyte/s", "tbyte/s"], expected_data={
            "byte/s (inc)": [1.953125e+11, 4.88281250e+11, 4.88281250e+10,
                             4.88281250e+10], "gbyte/s (inc)": [195.3125, 488.28125, 48.828125, 48.828125],
            "tbyte/s (inc)": [0.195312, 0.48828125, 0.04882812, 0.04882812]
        }, sample_file=cuda_example_file)


def test_flops_derivation():
    derivation_metrics_test(
        metrics=["flop8/s", "gflop8/s", "tflop8/s"],
        expected_data={
            "flop8/s (inc)": [3.417969e+14, 4.88281250e+14, 4.88281250e+13,
                              4.88281250e+14], "gflop8/s (inc)": [341796.875, 488281.25, 48828.125, 488281.25],
            "tflop8/s (inc)": [341.796875, 488.28125, 48.828125, 488.28125]
        },
        sample_file=cuda_example_file,
    )


def test_diff_profile():
    gf, derived_metrics = parse(["time/s"], triton_example_file)
    gf2, _ = parse(["time/s"], cuda_example_file)
    gf = apply_diff_profile(gf, derived_metrics, cuda_example_file, ["time/s"], None, None, 0.0)
    assert "time/s (inc)" in gf.dataframe.columns
