import NvRules


def get_identifier():
    return "RuleWithTables"


def get_name():
    return "Rule with Tables"


def get_description():
    return "Print tables from an NvRule"


class CustomMetricData:
    def __init__(self, name, value, subtype):
        self.name = name
        self.value = value
        self.subtype = subtype

    def __str__(self) -> str:
        return f"{self.name} is a {self.subtype} metric and has value {self.value}"


def apply(handle):
    ctx = NvRules.get_context(handle)
    fe = ctx.frontend()

    # Generate a rule message, to which we will add two tables
    msg_id = fe.message(
        NvRules.MsgType.OPTIMIZATION,
        "Tables are a great way to present relevant metrics"
        " of a rule result in a structured manner."
        " Tables can be generated by calling IFrontend.generate_table()."
        " Here are a couple of examples:",
    )

    # First table
    #
    # A simple table with columns of three different types (int, float, str).
    # Make sure to specify column names in the header and provide the data
    # as a list of rows.
    #
    header = ["ID", "Metric", "Value", "Unit"]
    data_by_row = [
        [1, "gpu__time_duration.sum", 3.42, "us"],
        [2, "sm__throughput.avg.pct_of_peak_sustained_elapsed", 5.28, "%"],
        [
            3,
            "gpu__compute_memory_throughput.avg.pct_of_peak_sustained_elapsed",
            27.45,
            "%",
        ],
    ]
    fe.generate_table(msg_id, header, data_by_row)

    # Second table
    #
    # In addition to int, float, and str types, you can also use custom objects
    # as data, as long as they implement a __str__ method.
    # Optionally, you can configure the table, e.g., to include a title or description,
    # sort the data, or adjust the relative column widths.
    # Please refer to NvRules.IFrontend.generate_table.__doc__ for more information.
    #
    compute_metric = CustomMetricData(
        "sm__issue_active.avg.pct_of_peak_sustained_elapsed",
        5.27,
        "Compute",
    )
    memory_metric_1 = CustomMetricData(
        "lts__t_sectors.avg.pct_of_peak_sustained_elapsed",
        13.21,
        "Memory",
    )
    memory_metric_2 = CustomMetricData(
        "dram__cycles_active.avg.pct_of_peak_sustained_elapsed",
        27.45,
        "Memory",
    )
    header = ["Metric", "Value", "Description"]
    data_by_row = [
        [compute_metric.name, compute_metric.value, compute_metric],
        [memory_metric_1.name, memory_metric_1.value, memory_metric_1],
        [memory_metric_2.name, memory_metric_2.value, memory_metric_2],
    ]
    config = {
        "title": "Metric Table",
        "description": "These metrics are relevant for the analysis",
        "sort_by": {
            "column": "Value",  # sort by second column
            "order": "descending",  # put largest value first
        },
        "per_column_configs": {
            "Value": {  # use a column name or index to configure individual columns
                "tooltip": "Value of the metric",  # tooltip for the header
                "style": {
                    "header": {"bold": True},
                    "data": {"italic": True},
                },
            },
            "Description": {
                "relative_width": 0.4,
            },
        },
    }
    fe.generate_table(msg_id, header, data_by_row, config)
