# Copyright (c) Advanced Micro Devices, Inc.
# SPDX-License-Identifier:  MIT

import csv
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional, Union

import numpy as np
import pandas as pd

from utils import schema
from utils.logger import console_debug, console_error, console_warning
from utils.parser import eval_metric

################################################
# Global vars
################################################
XMIN = 0.01
XMAX = 1000

FONT_SIZE = 16
FONT_COLOR = "black"
FONT_WEIGHT = "bold"

# SUPPORTED_DATATYPES table is based on datatype support in rocm-amdgpu-bench repository
# Indicates which datatypes per gpu arch can be generated by the roofline binary
SUPPORTED_DATATYPES: dict[str, list[str]] = {
    "gfx90a": [
        "FP16",
        "BF16",
        "FP32",
        "FP64",
        "I8",
        "I32",
        "I64",
    ],  # Unsupported: F4, F6, F8
    "gfx940": [
        "FP8",
        "FP16",
        "BF16",
        "FP32",
        "FP64",
        "I8",
        "I32",
        "I64",
    ],  # Unsupported: F4, F6
    "gfx941": [
        "FP8",
        "FP16",
        "BF16",
        "FP32",
        "FP64",
        "I8",
        "I32",
        "I64",
    ],  # Unsupported: F4, F6
    "gfx942": [
        "FP8",
        "FP16",
        "BF16",
        "FP32",
        "FP64",
        "I8",
        "I32",
        "I64",
    ],  # Unsupported: F4, F6
    "gfx950": [
        "FP4",
        "FP6",
        "FP8",
        "FP16",
        "BF16",
        "FP32",
        "FP64",
        "I8",
        "I32",
        "I64",
    ],
}

PEAK_OPS_DATATYPES = ["FP16", "FP32", "FP64", "I8", "I32", "I64"]
MATRIX_DATATYPES = ["FP4", "FP6", "FP8", "FP16", "BF16", "FP32", "FP64", "I8"]
CACHE_HIERARCHY = ["HBM", "L2", "L1", "LDS"]

TOP_N = 10


################################################
# Helper funcs
################################################
@dataclass
class AI_Data:
    KernelName: str
    numCalls: float

    total_flops: float
    valu_flops: float
    mfma_flops_f6f4: float
    mfma_flops_f8: float
    mfma_flops_f16: float
    mfma_flops_bf16: float
    mfma_flops_f32: float
    mfma_flops_f64: float
    mfma_iops_i8: float
    lds_data: float
    L1cache_data: float
    L2cache_data: float
    hbm_data: float

    totalDuration: float
    avgDuration: float


@dataclass
class PlotPoints:
    """Data structure for storing roofline plot points."""

    ai_l1: list[list[float]]
    ai_l2: list[list[float]]
    ai_hbm: list[list[float]]
    kernelNames: list[str]

    @classmethod
    def empty(cls) -> "PlotPoints":
        """Create empty plot points structure."""
        return cls(ai_l1=[[], []], ai_l2=[[], []], ai_hbm=[[], []], kernelNames=[])


@dataclass
class GraphPoints:
    """Data structure for storing roofline graph ceiling points."""

    hbm: list[Union[list[float], float, None]]
    l2: list[Union[list[float], float, None]]
    l1: list[Union[list[float], float, None]]
    lds: list[Union[list[float], float, None]]
    valu: list[Union[list[float], float, None]]
    matrix_ops: list[Union[list[float], float, None]]

    @classmethod
    def empty(cls) -> "GraphPoints":
        """Create empty graph points structure."""
        return cls(
            hbm=[None, None, None],
            l2=[None, None, None],
            l1=[None, None, None],
            lds=[None, None, None],
            valu=[None, None, None],
            matrix_ops=[None, None, None],
        )


################################################
# Helper functions
################################################
def get_font() -> dict[str, Union[int, str]]:
    return {
        "size": FONT_SIZE,
        "color": FONT_COLOR,
        "weight": FONT_WEIGHT,
        "family": "serif",
    }


def get_color(category: str) -> str:
    color_map = {"ai_l1": "green", "ai_l2": "blue", "ai_hbm": "red"}

    if category not in color_map:
        raise RuntimeError(f"Invalid category passed to get_color(): {category}")

    return color_map[category]


def sanitize_ai_value(value: float) -> float:
    excluded_values = ("", "N/A", np.inf, -np.inf, None)
    return value if value and value not in excluded_values else 0


# -------------------------------------------------------------------------------------
#                           Plot BW at each cache level
# -------------------------------------------------------------------------------------
def calc_ceilings(
    roofline_parameters: dict[str, Any],
    dtype: str,
    benchmark_data: dict[str, list[str]],
    ai_data: Optional[dict] = None,
) -> dict[str, list[Union[list[float], float, None]]]:
    """Given benchmarking data, calculate ceilings (or peak performance) for
    empirical roofline"""

    if ai_data:
        max_ai = 0
        for cache_level in ["ai_l1", "ai_l2", "ai_hbm"]:
            if cache_level in ai_data and ai_data[cache_level][0]:
                cache_max = max(ai_data[cache_level][0])
                max_ai = max(max_ai, cache_max)

        dynamic_xmax = max_ai * 1.2 if max_ai > 0 else 1000
    else:
        dynamic_xmax = 1000

    # TODO: This is where filtering by memory level will need to occur for standalone
    graph_points: dict[str, list[Union[list[float], float, None]]] = {
        "hbm": [],
        "l2": [],
        "l1": [],
        "lds": [],
        "valu": [],
        "matrix_ops": [],
    }

    mem_level = roofline_parameters["mem_level"]
    cache_hierarchy = (
        CACHE_HIERARCHY if mem_level == "ALL" or mem_level == ["ALL"] else mem_level
    )

    x1 = y1 = x2 = y2 = -1
    x1_matrix = y1_matrix = x2_matrix = y2_matrix = -1

    ops_flops = "Ops" if dtype.startswith("I") else "Flops"

    peak_ops = 0.0
    if dtype in PEAK_OPS_DATATYPES:
        try:
            peak_ops = float(
                benchmark_data[f"{dtype}{ops_flops}"][roofline_parameters["device_id"]]
            )
        except KeyError:
            console_warning(
                f"Missing benchmark data for {dtype}{ops_flops} in benchmark_results. "
                "Skipping peak operations calculation. This may indicate incomplete or "
                "corrupted benchmark data."
            )
            return GraphPoints.empty().__dict__

    for cache_level in cache_hierarchy:
        # Plot BW line
        curr_bw = f"{cache_level}Bw"
        try:
            peak_bw = float(benchmark_data[curr_bw][roofline_parameters["device_id"]])
        except KeyError:
            console_warning(
                f"Missing benchmark data for {curr_bw} in benchmark_results. "
                f"Skipping {cache_level} cache level. This may indicate incomplete or "
                "corrupted benchmark data."
            )
            continue

        x1 = float(XMIN)
        y1 = float(XMIN) * peak_bw

        if dtype in PEAK_OPS_DATATYPES:
            x2 = peak_ops / peak_bw
            y2 = peak_ops  # noqa

            # Plot Matrix Ops lines (NOTE: Assuming MI200 soc)
            x1_matrix = peak_ops / peak_bw
            y1_matrix = peak_ops

        peak_matrix = 0.0
        if dtype in MATRIX_DATATYPES:
            target_precision = dtype if dtype.startswith("I") else f"F{dtype[2:]}"

            try:
                peak_matrix = float(
                    benchmark_data[f"MFMA{target_precision}{ops_flops}"][
                        roofline_parameters["device_id"]
                    ]
                )
                x2_matrix = peak_matrix / peak_bw
                y2_matrix = peak_matrix
            except KeyError:
                console_warning(
                    f"Missing benchmark data for "
                    f"MFMA{target_precision}{ops_flops} in benchmark_results. "
                    f"Skipping MFMA calculations for {cache_level} cache level. "
                    "This may indicate incomplete or corrupted benchmark data."
                )

        # Check which peak is higher for formatting bandwidth lines
        if y2_matrix > y1_matrix:  # peak_matrix
            peak_x = x2_matrix
            peak_y = y2_matrix
        else:  # peakVALU
            peak_x = x1_matrix
            peak_y = y1_matrix

        cache_key = cache_level.lower()
        graph_points[cache_key].extend([[x1, peak_x], [y1, peak_y], peak_bw])

    # ----------------------------------------------------------------------------------
    #                                     Plot computing roof
    # ----------------------------------------------------------------------------------
    if dtype in PEAK_OPS_DATATYPES:
        # Plot FMA roof
        x0 = min(x2, dynamic_xmax) if x2 < dynamic_xmax else dynamic_xmax

        graph_points["valu"].extend([
            [x0, dynamic_xmax],
            [peak_ops, peak_ops],
            peak_ops,
        ])

    # Plot Matrix Ops roof
    if dtype in MATRIX_DATATYPES:  # assert that "matrix_ops" has been assigned
        x0_matrix = (
            min(x2_matrix, dynamic_xmax) if x2_matrix < dynamic_xmax else dynamic_xmax
        )

        graph_points["matrix_ops"].extend([
            [x0_matrix, dynamic_xmax],
            [peak_matrix, peak_matrix],
            peak_matrix,
        ])

    return graph_points


# -------------------------------------------------------------------------------------
#                              Overlay application performance
# -------------------------------------------------------------------------------------
# Calculate relevant metrics for ai calculation
def calc_ai_analyze(
    workload: schema.Workload,
    pmc_df: pd.DataFrame,
    config: dict[str, Any],
    arch_config: schema.ArchConfig,
) -> dict[str, Union[list[list[float]], list[str]]]:
    """
    Calculate per-kernel metrics and AI points with Roofline yamls using eval_metric.

    Caller must apply filters (e.g. parser.apply_filters) and pass the resulting
    DataFrame as pmc_df. This function only computes AI data points.
    """
    console_debug("calc_ai_analyze", "Starting calc_ai analysis using Roofline yamls")
    plot_points = PlotPoints.empty()

    workload.roofline_metrics = {}

    kernel_ids_to_process: list[int] = []
    kernel_top_table_id = 1

    if workload.filter_kernel_ids:
        if all(isinstance(k, int) for k in workload.filter_kernel_ids):
            kernel_ids_to_process = workload.filter_kernel_ids
        elif kernel_top_table_id in workload.dfs:
            # WebUI sets filter_kernel_ids as strings (kernel names), not
            # int row indices. Process all IDs here; caller provides
            # pmc_df already narrowed to the selected kernel(s).
            kernel_ids_to_process = workload.dfs[kernel_top_table_id].index.tolist()
    elif kernel_top_table_id in workload.dfs:
        kernel_top_df = workload.dfs[kernel_top_table_id]
        kernel_ids_to_process = kernel_top_df.index.tolist()
        console_debug(
            "roofline", f"Found {len(kernel_ids_to_process)} kernels to process"
        )

    if not kernel_ids_to_process:
        console_warning("No kernels found to process for roofline")
        return plot_points.__dict__

    for kernel_id in kernel_ids_to_process:
        kernel_name = ""
        if kernel_top_table_id in workload.dfs:
            kernel_top_df = workload.dfs[kernel_top_table_id]
            if kernel_id not in kernel_top_df.index:
                continue
            kernel_name = kernel_top_df.loc[kernel_id, "Kernel_Name"]
        else:
            continue

        console_debug("roofline", f"Processing kernel {kernel_id}: {kernel_name[:50]}")

        # filter PMC data for specific kernel
        kernel_pmc_df = pmc_df[pmc_df["pmc_perf"]["Kernel_Name"] == kernel_name]

        if kernel_pmc_df.empty:
            console_debug("roofline", f"No PMC data for kernel {kernel_id}")
            continue

        kernel_only_data = {"pmc_perf": kernel_pmc_df["pmc_perf"]}

        kernel_dfs: dict[int, pd.DataFrame] = {}
        kernel_dfs_type: dict[int, str] = {}

        for table_id in [401, 402]:
            if table_id in arch_config.dfs:
                kernel_dfs[table_id] = arch_config.dfs[table_id].copy()
                kernel_dfs_type[table_id] = arch_config.dfs_type[table_id]

        # eval metrics for single kernel only
        eval_metric(
            kernel_dfs,
            kernel_dfs_type,
            workload.sys_info.iloc[0],
            workload.roofline_peaks,
            kernel_only_data,
            debug=False,
            config=config,
        )

        ai_hbm = ai_l2 = ai_l1 = performance = 0

        if 402 in kernel_dfs:
            for idx, row in kernel_dfs[402].iterrows():
                metric = row.get("Metric", "")
                value = row.get("Value", 0)
                if metric == "AI HBM":
                    ai_hbm = sanitize_ai_value(value)
                elif metric == "AI L2":
                    ai_l2 = sanitize_ai_value(value)
                elif metric == "AI L1":
                    ai_l1 = sanitize_ai_value(value)
                elif metric == "Performance (GFLOPs)":
                    performance = sanitize_ai_value(value)

        console_debug(
            "roofline",
            f"Kernel {kernel_id}: "
            f"AI_HBM={ai_hbm:.2f}, "
            f"AI_L2={ai_l2:.2f}, "
            f"AI_L1={ai_l1:.2f}, "
            f"Performance={performance:.2e} GFLOP/s",
        )

        # add to plot points if we have valid data
        if performance > 0:
            if ai_hbm >= 0:
                plot_points.ai_hbm[0].append(ai_hbm)
                plot_points.ai_hbm[1].append(performance)
            if ai_l2 >= 0:
                plot_points.ai_l2[0].append(ai_l2)
                plot_points.ai_l2[1].append(performance)
            if ai_l1 >= 0:
                plot_points.ai_l1[0].append(ai_l1)
                plot_points.ai_l1[1].append(performance)

            plot_points.kernelNames.append(f"K{kernel_id}")
            console_debug("roofline", f"Added kernel {kernel_id} to plot points")
        else:
            console_debug(
                "roofline", f"Skipping kernel {kernel_id} - no performance data"
            )

        # store metrics for display
        workload.roofline_metrics[kernel_id] = {
            "name": kernel_name,
            "ai_table": kernel_dfs.get(401, pd.DataFrame()),
            "calc_table": kernel_dfs.get(402, pd.DataFrame()),
        }

    console_debug("roofline", f"Generated {len(plot_points.kernelNames)} plot points")
    return plot_points.__dict__


def construct_roof(
    roofline_parameters: dict[str, Any], dtype: str, ai_data: Optional[dict] = None
) -> dict[str, list[Union[list[float], float, None]]]:
    workload_dir = roofline_parameters.get("workload_dir")

    # Normalize workload_dir to extract base directory
    if isinstance(workload_dir, list):
        base_dir = (
            workload_dir[0][0]
            if isinstance(workload_dir[0], (list, tuple))
            else workload_dir[0]
        )
    else:
        base_dir = workload_dir

    benchmark_results = Path(base_dir) / "roofline.csv"

    # Initialize benchmark data dictionary from roofline.csv
    benchmark_data: dict[str, list[str]] = {}
    headers: list[str] = []

    try:
        with open(benchmark_results) as csvfile:
            csv_reader = csv.reader(csvfile, delimiter=",")
            row_count = 0

            for row in csv_reader:
                row.pop(0)  # Remove first column (Device ID)
                if row_count == 0:
                    headers = row
                    for header in headers:
                        benchmark_data[header] = []
                else:
                    for i, key in enumerate(headers):
                        benchmark_data[key].append(row[i])
                row_count += 1
    except Exception as e:
        console_error(
            "roofline",
            f"Failed to read benchmark results from {base_dir}: {e}",
            exit=False,
        )
        return GraphPoints.empty().__dict__

    # ------------------
    #  Validate benchmark data completeness
    # ------------------
    ops_flops = "Ops" if dtype.startswith("I") else "Flops"
    expected_columns = []

    if dtype in PEAK_OPS_DATATYPES:
        expected_columns.append(f"{dtype}{ops_flops}")

    mem_level = roofline_parameters["mem_level"]
    cache_hierarchy = (
        CACHE_HIERARCHY if mem_level == "ALL" or mem_level == ["ALL"] else mem_level
    )
    for cache_level in cache_hierarchy:
        expected_columns.append(f"{cache_level}Bw")

    if dtype in MATRIX_DATATYPES:
        target_precision = dtype if dtype.startswith("I") else f"F{dtype[2:]}"
        expected_columns.append(f"MFMA{target_precision}{ops_flops}")

    # Check for missing expected columns
    missing_columns = [col for col in expected_columns if col not in benchmark_data]
    if missing_columns:
        console_warning(
            f"Missing expected columns in roofline.csv for datatype {dtype}: "
            f"{', '.join(missing_columns)}. "
            "The roofline plot may be incomplete. Consider regenerating "
            "benchmark data or cleaning the directory and re-running the analysis."
        )

    # ------------------
    #  Generate Roofline
    # ------------------
    return calc_ceilings(roofline_parameters, dtype, benchmark_data, ai_data)
