plot.py

"""Generate MangoHud and Phoronix FPS charts from CSV exports.

This script scans a directory for CSV files and produces:
1. A MangoHud FPS distribution heatmap across runs.
2. A Phoronix average FPS bar chart per export CSV.

For Phoronix exports that do not include run statistics, the script can
enrich run count, variance, and standard error using matching MangoHud runs.
"""

import glob
import os
import csv
import re
import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.ticker import EngFormatter

FPS_SUBDIVS = 1.0
MAX_PRECISE_SAMPLES = 10000


def percentile_from_sorted(sorted_vals, pct):
    # Match MangoPlot's percentile-by-index behavior on sorted FPS values.
    if len(sorted_vals) == 0:
        return np.nan
    idx = int(len(sorted_vals) * (pct / 100.0))
    idx = max(0, min(len(sorted_vals) - 1, idx))
    return float(sorted_vals[idx])


def trim_trailing_frozen_tail(fps_values, ft_values, min_tail_len=30):
    # Some captures end with repeated frozen frames after the run; trim that trailing plateau.
    n = len(fps_values)
    if n == 0:
        return fps_values, ft_values, 0
    last_fps = fps_values[-1]
    last_ft = ft_values[-1]
    tail_len = 0
    for i in range(n - 1, -1, -1):
        if np.isclose(fps_values[i], last_fps, rtol=0.0, atol=1e-6) and np.isclose(ft_values[i], last_ft, rtol=0.0, atol=1e-6):
            tail_len += 1
        else:
            break

    if tail_len >= min_tail_len and tail_len < n:
        keep_until = n - tail_len
        return fps_values[:keep_until], ft_values[:keep_until], tail_len

    return fps_values, ft_values, 0


def robust_percentile(sorted_vals, pct):
    # Keep MangoPlot percentile-index behavior even for shorter captures.
    return percentile_from_sorted(sorted_vals, pct)


def get_csv_files(log_dir):
    return [f for f in glob.glob(os.path.join(log_dir, "*.csv")) if not f.endswith("_summary.csv")]


def coerce_float(value):
    value_num = pd.to_numeric(pd.Series([value]), errors="coerce").iloc[0]
    return float(value_num) if pd.notna(value_num) else None


def coerce_int(value):
    value_num = coerce_float(value)
    return int(value_num) if value_num is not None else None


def find_header_row(filepath, max_scan_rows=100):
    # Match MangoPlot's behavior of scanning for the row containing the fps column.
    with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
        for row_idx, line in enumerate(f):
            if row_idx > max_scan_rows:
                break
            cols = [c.strip() for c in line.strip().split(",")]
            if "fps" in cols:
                return row_idx
    return None


def parse_mangohud_files(csv_files):
    records = []
    for f in sorted(csv_files):
        header_row = find_header_row(f)
        if header_row is None:
            continue

        df = pd.read_csv(f, skiprows=header_row)
        if "fps" not in df.columns or "frametime" not in df.columns:
            continue
        fps = pd.to_numeric(df["fps"], errors="coerce")
        ft = pd.to_numeric(df["frametime"], errors="coerce")
        valid = fps.notna() & ft.notna() & (fps > 0.0) & (ft > 0.0)
        fps_values = fps[valid].to_numpy()
        ft_values = ft[valid].to_numpy()
        if len(fps_values) == 0:
            continue
        fps_values, ft_values, trimmed_tail = trim_trailing_frozen_tail(fps_values, ft_values)
        if len(fps_values) == 0:
            continue
        fps_sorted = np.sort(fps_values)

        bar_distribution = []
        for fps_val in fps_sorted:
            if fps_val > 1000:
                continue
            index = int(fps_val / FPS_SUBDIVS)
            if index >= len(bar_distribution):
                bar_distribution.extend([0] * (index + 1 - len(bar_distribution)))
            bar_distribution[index] += 1

        if not bar_distribution:
            continue

        records.append({
            "run": os.path.basename(f).replace(".csv", ""),
            "low_0_1pct": robust_percentile(fps_sorted, 0.1),
            "low_1pct": robust_percentile(fps_sorted, 1.0),
            "p50_fps": robust_percentile(fps_sorted, 50.0),
            "avg_fps": float(np.mean(fps_values)),
            "avg_frametime": float(np.mean(ft_values)),
            "p99_frametime": float(np.percentile(ft_values, 99)),
            "samples": len(fps_values),
            "trimmed_tail": trimmed_tail,
            "is_precise": len(fps_values) >= MAX_PRECISE_SAMPLES,
            "distribution": bar_distribution,
        })
    return records


def normalize_text(value):
    return re.sub(r"[^a-z0-9]+", "", value.lower())


def mangohud_group_key(run_name):
    # Strip trailing timestamp chunks from names like Game_YYYY-MM-DD_HH-MM-SS.
    return re.sub(r"_[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}$", "", run_name)


def build_mangohud_fallback_stats(records):
    grouped = {}
    for record in records:
        key = mangohud_group_key(record["run"])
        grouped.setdefault(key, []).append(float(record["avg_fps"]))

    fallback = {}
    for key, avg_fps_values in grouped.items():
        if len(avg_fps_values) == 0:
            continue

        arr = np.array(avg_fps_values, dtype=float)
        runs = len(arr)
        variance = float(np.var(arr, ddof=1)) if runs > 1 else 0.0
        stderr = float(np.std(arr, ddof=1) / np.sqrt(runs)) if runs > 1 else 0.0
        fallback[key] = {
            "runs": runs,
            "variance": variance,
            "stderr": stderr,
            "normalized": normalize_text(key),
        }

    return fallback


def match_mangohud_stats_for_benchmark(benchmark_name, fallback_stats):
    benchmark_norm = normalize_text(benchmark_name)
    best_key = None
    best_len = -1

    for key, stats in fallback_stats.items():
        key_norm = stats["normalized"]
        if not key_norm:
            continue
        if key_norm in benchmark_norm or benchmark_norm in key_norm:
            if len(key_norm) > best_len:
                best_key = key
                best_len = len(key_norm)

    if best_key is None:
        return None
    return fallback_stats[best_key]


def enrich_phoronix_with_mangohud(benchmarks, mangohud_records):
    if not benchmarks or not mangohud_records:
        return benchmarks

    fallback_stats = build_mangohud_fallback_stats(mangohud_records)
    if not fallback_stats:
        return benchmarks

    enriched = []
    for benchmark in benchmarks:
        updated = dict(benchmark)
        match = match_mangohud_stats_for_benchmark(updated.get("benchmark", ""), fallback_stats)
        if match is None:
            enriched.append(updated)
            continue

        current_runs = pd.to_numeric(pd.Series([updated.get("runs")]), errors="coerce").iloc[0]
        # PTS CSV exports often only provide one aggregate score row (HIB), which infers N=1.
        # If MangoHud has grouped runs, prefer that richer run count.
        if pd.isna(current_runs) or (int(current_runs) <= 1 and int(match["runs"]) > 1):
            updated["runs"] = int(match["runs"])

        current_variance = pd.to_numeric(pd.Series([updated.get("variance")]), errors="coerce").iloc[0]
        if pd.isna(current_variance):
            updated["variance"] = float(match["variance"])

        current_stderr = pd.to_numeric(pd.Series([updated.get("stderr")]), errors="coerce").iloc[0]
        if pd.isna(current_stderr):
            updated["stderr"] = float(match["stderr"])

        enriched.append(updated)

    return enriched


def parse_phoronix_csv(filepath):
    benchmark_stats = {}
    current_system = ""
    metadata_rows = {
        "processor",
        "motherboard",
        "chipset",
        "memory",
        "disk",
        "graphics",
        "audio",
        "monitor",
        "network",
        "os",
        "kernel",
        "display server",
        "opengl",
        "opencl",
        "compiler",
        "file-system",
        "screen resolution",
    }

    with open(filepath, "r", encoding="utf-8", errors="ignore", newline="") as csv_file:
        reader = csv.reader(csv_file)
        for row in reader:
            row = [cell.strip() for cell in row]
            if len(row) == 0:
                continue
            if len(row) >= 3 and not row[0] and not row[1] and row[2]:
                current_system = row[2]
                continue

            label = row[0]
            if label.lower() in metadata_rows:
                continue

            # Keep only rows with numeric score-like values in the third column.
            try:
                score = float(row[2])
            except (ValueError, IndexError):
                continue

            result_type = row[1].strip().upper() if len(row) > 1 else ""
            stats = benchmark_stats.setdefault(
                label,
                {
                    "benchmark": label,
                    "avg_fps": None,
                    "system": current_system or os.path.basename(filepath).replace(".csv", ""),
                    "runs": None,
                    "variance": None,
                    "raw_scores": [],
                    "stddev": None,
                    "stderr": None,
                },
            )

            if result_type in {"HIB", "LIB"}:
                stats["avg_fps"] = score
                stats["raw_scores"].append(score)
                continue

            if "RUN" in result_type or result_type == "N":
                stats["runs"] = int(score)
                continue

            if "VAR" in result_type:
                stats["variance"] = score
                continue

            if "DEV" in result_type or result_type in {"SD", "STD"}:
                stats["stddev"] = score
                continue

            if "ERR" in result_type or result_type == "SE":
                stats["stderr"] = score
                continue

    benchmarks = []
    for stats in benchmark_stats.values():
        if stats["avg_fps"] is None:
            continue

        if stats["runs"] is None:
            stats["runs"] = len(stats["raw_scores"]) if len(stats["raw_scores"]) > 0 else 1

        if stats["variance"] is None:
            if stats["stddev"] is not None:
                stats["variance"] = float(stats["stddev"]) ** 2
            elif len(stats["raw_scores"]) > 1:
                stats["variance"] = float(np.var(stats["raw_scores"], ddof=1))

        benchmarks.append({
            "benchmark": stats["benchmark"],
            "avg_fps": stats["avg_fps"],
            "system": stats["system"],
            "runs": stats["runs"],
            "variance": stats["variance"],
            "stderr": stats["stderr"],
        })

    return benchmarks


def save_phoronix_avg_fps_chart(benchmarks, source_csv):
    if not benchmarks:
        return None

    df = pd.DataFrame(benchmarks)
    df = df.sort_values("avg_fps", ascending=False).reset_index(drop=True)
    chart_height = max(3.0, 1.0 + 0.7 * len(df))
    fig, ax = plt.subplots(figsize=(12, chart_height))
    fig.patch.set_facecolor("#1A1C1D")
    ax.set_facecolor("#1A1C1D")

    y_positions = np.arange(len(df))
    bars = ax.barh(y_positions, df["avg_fps"], color="#0967BA", edgecolor="#e8e6e3", linewidth=0.8, height=0.58)
    ax.invert_yaxis()

    max_fps = float(df["avg_fps"].max())
    x_max = max_fps * 1.18 if max_fps > 0 else 1.0
    ax.set_xlim(0.0, x_max)

    ax.set_yticks(y_positions, labels=df["benchmark"])
    ax.tick_params(axis="x", colors="#e8e6e3")
    ax.tick_params(axis="y", colors="#e8e6e3")
    ax.xaxis.grid(True, color="#585f63", linestyle=(0, (5, 5)), linewidth=1.0)
    ax.set_axisbelow(True)
    ax.set_xlabel("Frames Per Second, More Is Better", color="#e8e6e3", fontweight="bold")

    title = "Average FPS (Phoronix CSV Export)"
    system_name = str(df["system"].iloc[0]).strip()
    if system_name:
        title += f" - {system_name}"
    ax.set_title(title, color="#e8e6e3", fontweight="bold", fontsize=13, loc="left")

    for idx, (bar, value) in enumerate(zip(bars, df["avg_fps"])):
        y = bar.get_y() + bar.get_height() / 2.0
        width = bar.get_width()
        runs = df.loc[idx, "runs"]
        variance = df.loc[idx, "variance"]
        stderr = df.loc[idx, "stderr"]

        value_label = f"{value:.2f}"
        variance_label = f"{variance:.4f}" if pd.notna(variance) else "n/a"
        runs_num = coerce_int(runs)
        runs_label = str(runs_num) if runs_num is not None else "n/a"
        stats_label = f"N={runs_label}, Var={variance_label}"
        stderr_value = coerce_float(stderr)
        if stderr_value is not None:
            stats_label += f", SE={stderr_value:.2f}"

        if width > x_max * 0.15:
            ax.text(width - (x_max * 0.01), y, value_label, va="center", ha="right", color="#FFFFFF", fontweight="bold", fontsize=10)
            ax.text(width - (x_max * 0.01), y + 0.24, stats_label, va="center", ha="right", color="#e8e6e3", fontsize=8)
        else:
            ax.text(width + (x_max * 0.01), y, value_label, va="center", ha="left", color="#e8e6e3", fontweight="bold", fontsize=10)
            ax.text(width + (x_max * 0.01), y + 0.24, stats_label, va="center", ha="left", color="#e8e6e3", fontsize=8)

    for spine in ["top", "right"]:
        ax.spines[spine].set_visible(False)
    ax.spines["left"].set_color("#e8e6e3")
    ax.spines["bottom"].set_color("#e8e6e3")

    plt.tight_layout()
    output_path = os.path.join(os.path.dirname(source_csv), f"{os.path.basename(source_csv).replace('.csv', '')}_avg_fps_bar_chart.png")
    plt.savefig(output_path, dpi=160, transparent=True)
    plt.close(fig)
    return output_path


def save_mangohud_chart(records, output_dir):
    if not records:
        return None

    records = sorted(records, key=lambda r: r["run"])
    out = pd.DataFrame([{k: v for k, v in r.items() if k != "distribution"} for r in records])
    if out.empty:
        return None

    distributions = [r["distribution"] for r in records]
    y_labels = [r["run"] for r in records]
    if not distributions:
        return None

    if not out["is_precise"].all():
        for run_name in out.loc[~out["is_precise"], "run"].tolist():
            print(f"Warning: '{run_name}' has < {MAX_PRECISE_SAMPLES} samples; low-percentile stats are less precise.")

    print(
        out[
            [
                "run",
                "low_0_1pct",
                "low_1pct",
                "p50_fps",
                "avg_fps",
                "avg_frametime",
                "p99_frametime",
                "samples",
                "trimmed_tail",
                "is_precise",
            ]
        ].to_string(index=False)
    )

    max_size = max(len(d) for d in distributions)
    for d in distributions:
        if len(d) < max_size:
            d.extend([0] * (max_size - len(d)))

    heatmap = np.array(distributions)
    num_runs = len(distributions)

    # Cap the x-axis to practical gameplay FPS stats, not rare spike outliers.
    max_stat_fps = float(out[["low_0_1pct", "low_1pct", "p50_fps", "avg_fps"]].to_numpy().max())
    x_headroom = 20.0
    display_max_fps = np.ceil(max_stat_fps / 10.0) * 10.0 + x_headroom
    display_bin_count = int(display_max_fps / FPS_SUBDIVS) + 1
    display_bin_count = min(display_bin_count, heatmap.shape[1])
    heatmap = heatmap[:, :display_bin_count]
    display_max_fps = float((display_bin_count - 1) * FPS_SUBDIVS)

    fig, ax = plt.subplots(figsize=(13, 7))
    fig.patch.set_facecolor("#1A1C1D")
    ax.set_facecolor("#1A1C1D")

    im = ax.imshow(
        heatmap,
        aspect="auto",
        extent=(0.0, display_max_fps, 0.0, float(num_runs)),
        cmap="magma",
    )

    for i in range(num_runs + 1):
        ax.axhline(float(i), color="#e8e6e3", lw=1.2, alpha=0.8)

    for i, row in enumerate(out.reset_index(drop=True).to_dict("records")):
        kwargs = dict(
            ymin=(num_runs - i - 1 + 0.15) / num_runs,
            ymax=(num_runs - i - 0.15) / num_runs,
            lw=2.5,
        )
        ax.axvline(float(row["low_0_1pct"]), color="#35260f", label=("0.1%" if i == 0 else None), **kwargs)
        ax.axvline(float(row["low_1pct"]), color="#6E4503", label=("1%" if i == 0 else None), **kwargs)
        ax.axvline(float(row["p50_fps"]), color="#0967BA", label=("50%" if i == 0 else None), **kwargs)

    for spine in ["left", "right", "bottom", "top"]:
        ax.spines[spine].set_color("#e8e6e3")
        ax.spines[spine].set_linewidth(1.2)

    ax.tick_params(axis="y", colors="#e8e6e3")
    ax.tick_params(axis="x", colors="#e8e6e3")
    ax.set_yticks(np.arange(num_runs - 0.5, 0, -1), labels=y_labels)
    ax.xaxis.set_major_formatter(EngFormatter(unit="FPS"))
    ax.set_xlabel("FPS", color="#e8e6e3")
    ax.set_title("FPS Distribution Heatmap (MangoPlot-style)", color="#e8e6e3")
    ax.grid(False)

    cbar = plt.colorbar(im, ax=ax, fraction=0.03, pad=0.02)
    cbar.set_label("Frame count", color="#e8e6e3")
    cbar.ax.yaxis.set_tick_params(color="#e8e6e3")
    plt.setp(cbar.ax.get_yticklabels(), color="#e8e6e3")

    legend = ax.legend(
        loc="upper left",
        bbox_to_anchor=(0.0, 1.08),
        ncol=2,
        facecolor="#585f63",
        edgecolor="#585f63",
    )
    for text in legend.get_texts():
        text.set_color("#cccbc9")

    plt.tight_layout()
    output_path = os.path.join(output_dir, "mangohud_bar_chart.png")
    plt.savefig(output_path, dpi=160, transparent=True)
    plt.close(fig)
    return output_path


def run(log_dir):
    files = get_csv_files(log_dir)

    mangohud_records = parse_mangohud_files(files)
    mangohud_output = save_mangohud_chart(mangohud_records, log_dir)
    if mangohud_output:
        print(f"\nSaved: {mangohud_output}")

    phoronix_outputs = []
    for csv_path in sorted(files):
        benchmarks = parse_phoronix_csv(csv_path)
        benchmarks = enrich_phoronix_with_mangohud(benchmarks, mangohud_records)
        output = save_phoronix_avg_fps_chart(benchmarks, csv_path)
        if output:
            phoronix_outputs.append(output)

    for output_path in phoronix_outputs:
        print(f"Saved: {output_path}")

    if not mangohud_output and not phoronix_outputs:
        raise SystemExit("No usable MangoHud or Phoronix CSV files found.")


def main():
    parser = argparse.ArgumentParser(
        description="Generate MangoHud and PTS FPS charts from CSV logs.",
        epilog=(
            "Examples:\n"
            "  python plot.py\n"
            "  python plot.py --log-dir /path/to/benchmark-logs"
        ),
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument(
        "--log-dir",
        default=os.getcwd(),
        help="Directory containing CSV logs. Defaults to current working directory.",
    )
    args = parser.parse_args()
    run(os.path.abspath(args.log_dir))


if __name__ == "__main__":
    main()