#!/usr/bin/env python3


#
# Copyright(C) Advanced Micro Devices, Inc. All rights reserved.
#
# You may not use this software and documentation (if any) (collectively,
# the "Materials") except in compliance with the terms and conditions of
# the Software License Agreement included with the Materials or otherwise as
# set forth in writing and signed by you and an authorized signatory of AMD.
# If you do not have a copy of the Software License Agreement, contact your
# AMD representative for a copy.
#
# You agree that you will not reverse engineer or decompile the Materials,
# in whole or in part, except as allowed by applicable law.
#
# THE MATERIALS ARE DISTRIBUTED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR
# REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
#


"""
plot_json.py

--normalize    Scale each metric to 0-100 using (value / max * 100)
"""
import sys

__APP_VERSION__ = "1.130.9-a-3"
__PIPELINE__ = "rudra"
__P4_PROGRAM__ = "hydra"

if "-V" in sys.argv[1:] or "--version" in sys.argv[1:]:
    print(f"version: {__APP_VERSION__}")
    print(f"pipeline: {__PIPELINE__}")
    print(f"p4_program: {__P4_PROGRAM__}")
    sys.exit(0)

import argparse
import json
from pathlib import Path
import re # For sorting PBOQ, PAUSE, and Channel metrics

missing = []
try:
    import matplotlib.pyplot as plt
    from matplotlib.ticker import FuncFormatter # Import FuncFormatter
except ImportError as err:
    print(f"ImportError: {err}")
    missing.append("matplotlib")

try:
    import pandas as pd
except ImportError as err:
    print(f"ImportError: {err}")
    missing.append("pandas")

if missing:
    print(f"Error: Missing required packages: {', '.join(missing)}. Please install them and try again.")
    sys.exit(1)

# Define the metric profiles
METRIC_PROFILES = {
    "BW": [
        "PB_SWITCH/PBI/PBI_NET/PBI_NET_BPS",
        "PB_SWITCH/PBI/PBI_UD0/PBI_UD0_BPS",
        "PB_SWITCH/PBI/PBI_UD1/PBI_UD1_BPS",
        "PCI/PCI/BW_TXWR/BW_TxW",
        "PCI/PCI/BW_RXCMP/BW_RxC",
        # added PRD_BPS counters
        "UD0_P4DMA/PRD/PRD_BPS/PRD_BPS",
        "UD1_P4DMA/PRD/PRD_BPS/PRD_BPS",
    ],
    "PBOQ": [], # Placeholder, will be populated dynamically
    "RDWR": [
        "UD0_P4DMA/PRD/PRD Status/PRD_WRITES_CH0",
        "UD1_P4DMA/PRD/PRD Status/PRD_WRITES_CH0",
        "UD0_P4DMA/PTD/PTD Status/PT0_PTD_READS_CH0",
        "UD0_P4DMA/PTD/PTD Status/PT0_PTD_READS_CH1",
        "UD0_P4DMA/PTD/PTD Status/PT0_PTD_WRITES_CH0",
        "UD0_P4DMA/PTD/PTD Status/PT0_PTD_WRITES_CH1",
        "UD1_P4DMA/PTD/PTD Status/PT1_PTD_READS_CH0",
        "UD1_P4DMA/PTD/PTD Status/PT1_PTD_READS_CH1",
        "UD1_P4DMA/PTD/PTD Status/PT1_PTD_WRITES_CH0",
        "UD1_P4DMA/PTD/PTD Status/PT1_PTD_WRITES_CH1",
        "PCI/PCI/OUTSTANDING/READS",
    ],
    # new PHB profile: all 16 COS depths for UD0 and UD1
    "PHB": [
        *[f"UD0_P4DMA/PHB/QDEPTH/COS{i}" for i in range(16)],
        *[f"UD1_P4DMA/PHB/QDEPTH/COS{i}" for i in range(16)],
    ],
    # new PPS profile
    "PPS": [
        "UD0_P4DMA/PRD/PPS/PR_PPS",
        "UD1_P4DMA/PRD/PPS/PR_PPS",
        "UD0_P4DMA/PTD/PPS/PT0_PPS",
        "UD1_P4DMA/PTD/PPS/PT1_PPS",
    ],
    "PAUSE": [], # Placeholder, will be populated dynamically
    "RXFIFO": [
        "UD0_P4DMA/PRD/DFIFO/PKTFIFO",
        "UD1_P4DMA/PRD/DFIFO/PKTFIFO",
    ],
    "LATENCY": [], # Populated dynamically (PCI_RLATENCY/PCI_WLATENCY labels are hardware-dependent)
    "_LATENCY_STATIC": [
        "UD0_P4DMA/SU/AXI_LAT/AXI_LAT_NS",
        "UD1_P4DMA/SU/AXI_LAT/AXI_LAT_NS",
    ],
    "UD0CH": [], # Placeholder, will be populated dynamically
    "UD1CH": [], # Placeholder, will be populated dynamically
}


def flatten_iter(iteration):
    row = {"timestamp_ms": iteration["timestamp_ms"]}
    for pipe, p in iteration["pipelines"].items():
        for stage, s in p["stages"].items():
            for reg, r in s["registers"].items():
                for field, val in r.items():
                    row[f"{pipe}/{stage}/{reg}/{field}"] = val
    return row


def load_df(path: Path) -> pd.DataFrame:
    data = json.loads(path.read_text())
    rows = [flatten_iter(it) for it in data["iterations"]]
    df = pd.DataFrame(rows).set_index("timestamp_ms")
    df.index = df.index / 1_000.0
    df.index.name = "time_s"
    return df


def read_metric_file(fp: Path):
    return [
        ln.strip()
        for ln in fp.read_text().splitlines()
        if ln.strip() and not ln.lstrip().startswith("#")
    ]


def normalize_cols(df: pd.DataFrame) -> pd.DataFrame:
    max_vals = df.max().replace(0, 1)        # avoid /0
    return df.div(max_vals) * 100.0

# Custom formatter function for y-axis
def k_m_g_formatter(x, pos):
    if x == 0:
        return "0"
    abs_x = abs(x)
    if abs_x >= 1e9:
        return f'{x/1e9:.1f}G'
    elif abs_x >= 1e6:
        return f'{x/1e6:.1f}M'
    elif abs_x >= 1e3:
        return f'{x/1e3:.1f}K'
    else:
        # For smaller numbers, decide if decimals are needed or just int
        if x == int(x):
            return f'{int(x)}'
        else:
            return f'{x:.1f}'


def plot_metrics(df: pd.DataFrame, out: Path or None, normalized: bool):
    n_metrics = len(df.columns)
    if n_metrics > 20:
        fig, ax = plt.subplots(figsize=(16, 10))
        df.plot(ax=ax, lw=1.5, legend=False)
        ncols = 4 if n_metrics > 40 else 3
        ax.legend(fontsize=6, ncol=ncols, loc="upper center",
                  bbox_to_anchor=(0.5, -0.12))
    else:
        fig, ax = plt.subplots(figsize=(12, 6))
        df.plot(ax=ax, lw=1.5)
        ax.legend(fontsize="small")
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Value (0-100)" if normalized else "Value")
    ax.grid(True, alpha=0.3)

    if not normalized:
        ax.yaxis.set_major_formatter(FuncFormatter(k_m_g_formatter))

    if n_metrics > 20:
        plt.tight_layout(rect=[0, 0.15, 1, 1])
    else:
        plt.tight_layout()
    if out:
        if n_metrics > 20:
            plt.savefig(out, dpi=150, bbox_inches="tight")
        else:
            plt.savefig(out, dpi=150)
        print(f"✓ Saved to {out}")
    else:
        try:
            plt.show()
        except Exception as e:
            print(f"Warning: Could not display plot interactively ({e})")
            print("Use --out parameter to save to a file instead.")
            sys.exit(1)

def sort_pboq_metrics(metrics):
    """Sorts PBOQ metrics naturally by type (NET, UD0, UD1, AQ) and then index."""
    def get_sort_key(metric_name):
        parts = metric_name.split('/')
        field = parts[-1]
        match_field = re.match(r"PBI_(NET|UD0|UD1|AQ)_XOFF(\d+)", field)
        if match_field:
            type_prefix = match_field.group(1)
            idx = int(match_field.group(2))
            prefix_order = {"NET": 0, "UD0": 1, "UD1": 2, "AQ": 3}
            return (prefix_order.get(type_prefix, 99), idx, metric_name)
        return (999, 0, metric_name)
    return sorted(metrics, key=get_sort_key)

def sort_pause_metrics(metrics):
    """Sorts PAUSE metrics (e.g., PB_NET_XOFF[0-31]) by the number in brackets."""
    def get_sort_key(metric_name):
        match = re.search(r"(\d+)(?:\])?$", metric_name)
        if match:
            return (int(match.group(1)), metric_name)
        return (999, metric_name)
    return sorted(metrics, key=get_sort_key)

def sort_channel_metrics(metrics):
    """Sorts UDxCH metrics by Stage number, then Channel number."""
    def get_sort_key(metric_name):
        # Example: UD0_P4DMA/Stage 0/SDP_CHANNELS/CH0
        match = re.search(r"Stage (\d+)/SDP_CHANNELS/CH(\d+)", metric_name)
        if match:
            stage_num = int(match.group(1))
            channel_num = int(match.group(2))
            return (stage_num, channel_num, metric_name)
        return (99, 99, metric_name) # Fallback for non-matching names
    return sorted(metrics, key=get_sort_key)


def main():
    ap = argparse.ArgumentParser(description="Plot counters from pipeline JSON")
    ap.add_argument("-V", "--version", action="store_true",
                    help="display version")
    ap.add_argument("json_file", type=Path)
    ap.add_argument("--metrics", nargs="+", help="List of metrics to plot")
    ap.add_argument("--metrics-file", type=Path, help="File containing a list of metrics to plot")
    ap.add_argument("--profile", choices=[k for k in METRIC_PROFILES if not k.startswith("_")],
                    help="Name of a predefined metric profile to use")
    ap.add_argument("--list", action="store_true", help="List all available metrics from the JSON file and exit")
    ap.add_argument("--out", type=Path, help="Output file path to save the plot (e.g., plot.png)")
    ap.add_argument("--normalize", action="store_true",
                    help="Scale each metric so its max becomes 100")
    selector = ap.add_mutually_exclusive_group()
    selector.add_argument("-b", "--bdf", help="Select NIC by PCIe BDF")
    selector.add_argument("-c", "--card", help="Select NIC by UUID")
    args = ap.parse_args()

    df = load_df(args.json_file)

    if args.list:
        print("\n".join(sorted(df.columns)))
        sys.exit(0)

    chosen = []
    if args.profile:
        temp_chosen_metrics = [] # Common temp list for dynamic profiles
        if args.profile == "PBOQ":
            pboq_pattern = re.compile(r"/PBI_(NET|UD0|UD1|AQ)_XOFF(\d+)$")
            for col_name in df.columns:
                m = pboq_pattern.search(col_name)
                if m and 0 <= int(m.group(2)) <= 31:
                    temp_chosen_metrics.append(col_name)
            chosen.extend(sort_pboq_metrics(temp_chosen_metrics))
        elif args.profile == "LATENCY":
            for col_name in df.columns:
                if col_name.startswith("PCI/PCI/PCI_RLATENCY") or \
                   col_name.startswith("PCI/PCI/PCI_WLATENCY"):
                    temp_chosen_metrics.append(col_name)
            temp_chosen_metrics.sort()
            chosen.extend(temp_chosen_metrics)
            chosen.extend(METRIC_PROFILES["_LATENCY_STATIC"])
        elif args.profile == "PAUSE":
            pause_pattern = re.compile(r"/PBI_NET_XOFF(\d+)$")
            for col_name in df.columns:
                m = pause_pattern.search(col_name)
                if m and 0 <= int(m.group(1)) <= 31:
                    temp_chosen_metrics.append(col_name)
            chosen.extend(sort_pause_metrics(temp_chosen_metrics))
        elif args.profile == "UD0CH" or args.profile == "UD1CH":
            pipeline_prefix = "UD0_P4DMA" if args.profile == "UD0CH" else "UD1_P4DMA"
            # Regex to match: UDx_P4DMA/Stage <S>/SDP_CHANNELS/CH<C>
            # <S> is Stage number (0-7), <C> is Channel number (0-3)
            pattern = re.compile(f"^{re.escape(pipeline_prefix)}/Stage ([0-7])/SDP_CHANNELS/CH([0-3])$")
            for col_name in df.columns:
                match = pattern.match(col_name)
                if match:
                    # Stage and Channel numbers are already validated by the regex groups ([0-7] and [0-3])
                    temp_chosen_metrics.append(col_name)
            chosen.extend(sort_channel_metrics(temp_chosen_metrics))
        else: # For other static profiles
            chosen.extend(METRIC_PROFILES[args.profile])

    if args.metrics:
        chosen.extend(args.metrics)
    
    if args.metrics_file:
        chosen.extend(read_metric_file(args.metrics_file))
    
    seen = set()
    unique_chosen = []
    for x in chosen:
        if x not in seen:
            unique_chosen.append(x)
            seen.add(x)
    chosen = unique_chosen

    if not chosen:
        ap.error("Need metrics to plot. Use --metrics, --metrics-file, or --profile.")

    bad = [m for m in chosen if m not in df.columns]
    if bad:
        available_metrics_sample = "\n".join(sorted(df.columns)[:20]) 
        error_message = f"Unknown metric(s): {bad}\nUse --list to see all available metrics."
        if len(df.columns) > 20:
            error_message += f"\nFirst 20 available metrics:\n{available_metrics_sample}\n..."
        else:
            error_message += f"\nAvailable metrics:\n{available_metrics_sample}"
        ap.error(error_message)

    slice_df = df[chosen]
    if args.normalize:
        slice_df = normalize_cols(slice_df)

    plot_metrics(slice_df, args.out, args.normalize)


if __name__ == "__main__":
    main()
