# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "matplotlib",
#     "numpy",
#     "pandas",
# ]
# ///

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# NVIDIA GPUs only, sourced from original table
# cores from TechPowerUp specs
data_all = [
    {"time": 16, "gpu": "RTX4090", "cores": 16384, "vram": 24},
    {"time": 20, "gpu": "RTX4090", "cores": 16384, "vram": 24},
    {"time": 20, "gpu": "RTX3090", "cores": 10496, "vram": 24},
    {"time": 26, "gpu": "RTX4090", "cores": 16384, "vram": 24},
    {"time": 28, "gpu": "RTX4080 Super", "cores": 10240, "vram": 16},
    {"time": 30, "gpu": "RTX4070 Ti", "cores": 7680, "vram": 12},
    {"time": 30, "gpu": "RTX4070 Ti", "cores": 7680, "vram": 12},
    {"time": 30, "gpu": "RTX3080", "cores": 8704, "vram": 10},
    {"time": 30, "gpu": "RTX3080 Ti", "cores": 10240, "vram": 12},
    {"time": 32, "gpu": "RTX4080", "cores": 9728, "vram": 16},
    {"time": 32, "gpu": "RTX3080 Ti", "cores": 10240, "vram": 12},
    {"time": 32, "gpu": "RTX3080 12GB", "cores": 8960, "vram": 12},
    {"time": 36, "gpu": "RTX3080", "cores": 8704, "vram": 10},
    {"time": 36, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 36, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 36, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 36, "gpu": "RTX3070", "cores": 5888, "vram": 8},
    {"time": 38, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 40, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 42, "gpu": "RTX3070", "cores": 5888, "vram": 8},
    {"time": 42, "gpu": "RTX3070", "cores": 5888, "vram": 8},
    {"time": 44, "gpu": "RTX4060 Ti", "cores": 4352, "vram": 8},
    {"time": 44, "gpu": "RTX3060 Ti", "cores": 4864, "vram": 8},
    {"time": 46, "gpu": "RTX4060 Ti 16GB", "cores": 4608, "vram": 16},
    {"time": 48, "gpu": "RTX3060 Ti", "cores": 4864, "vram": 8},
    {"time": 50, "gpu": "RTX4060", "cores": 3072, "vram": 8},
    {"time": 50, "gpu": "RTX2080 Super", "cores": 3072, "vram": 8},
    {"time": 52, "gpu": "RTX4060", "cores": 3072, "vram": 8},
    {"time": 54, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 54, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 54, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 54, "gpu": "RTX4060", "cores": 3072, "vram": 8},
    {"time": 56, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 56, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 56, "gpu": "RTX4070 Mobile", "cores": 5888, "vram": 8},
    {"time": 58, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 58, "gpu": "RTX2070 Super", "cores": 2560, "vram": 8},
    {"time": 58, "gpu": "RTX4060 Mobile", "cores": 3072, "vram": 8},
    {"time": 58, "gpu": "RTX3060 Mobile", "cores": 3584, "vram": 6},
    {"time": 60, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 66, "gpu": "Quadro RTX4000", "cores": 2304, "vram": 8},
    {"time": 70, "gpu": "Quadro RTX4000", "cores": 2304, "vram": 8},
    {"time": 70, "gpu": "RTX3050 8GB", "cores": 2560, "vram": 8},
    {"time": 70, "gpu": "RTX2060", "cores": 1920, "vram": 6},
    {"time": 90, "gpu": "RTX-A1000", "cores": 1024, "vram": 4},
    {"time": 116, "gpu": "RTX3050 Ti Mobile", "cores": 2560, "vram": 4},
    {"time": 146, "gpu": "GTX1080", "cores": 2560, "vram": 8},
    {"time": 160, "gpu": "RTX3050 Ti Mobile", "cores": 2560, "vram": 4},
    {"time": 168, "gpu": "GTX1070 Ti", "cores": 2432, "vram": 8},
    {"time": 178, "gpu": "GTX1660 Ti", "cores": 1536, "vram": 6},
    {"time": 190, "gpu": "GTX1070", "cores": 1920, "vram": 8},
    {"time": 194, "gpu": "GTX1080", "cores": 2560, "vram": 8},
    {"time": 200, "gpu": "RTX3050 Ti Mobile", "cores": 2560, "vram": 4},
    {"time": 204, "gpu": "GTX1660", "cores": 1408, "vram": 6},
    {"time": 246, "gpu": "GTX1650", "cores": 896, "vram": 4},
    {"time": 284, "gpu": "GTX1060 6GB", "cores": 1280, "vram": 6},
]

data_filtered = [
    {"time": 28, "gpu": "RTX4080 Super", "cores": 10240, "vram": 16},
    {"time": 30, "gpu": "RTX4070 Ti", "cores": 7680, "vram": 12},
    {"time": 30, "gpu": "RTX4070 Ti", "cores": 7680, "vram": 12},
    {"time": 30, "gpu": "RTX3080", "cores": 8704, "vram": 10},
    {"time": 30, "gpu": "RTX3080 Ti", "cores": 10240, "vram": 12},
    {"time": 32, "gpu": "RTX4080", "cores": 9728, "vram": 16},
    {"time": 32, "gpu": "RTX3080 Ti", "cores": 10240, "vram": 12},
    {"time": 32, "gpu": "RTX3080 12GB", "cores": 8960, "vram": 12},
    {"time": 36, "gpu": "RTX3080", "cores": 8704, "vram": 10},
    {"time": 36, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 36, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 36, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 36, "gpu": "RTX3070", "cores": 5888, "vram": 8},
    {"time": 38, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 40, "gpu": "RTX4070", "cores": 5888, "vram": 12},
    {"time": 42, "gpu": "RTX3070", "cores": 5888, "vram": 8},
    {"time": 42, "gpu": "RTX3070", "cores": 5888, "vram": 8},
    {"time": 44, "gpu": "RTX4060 Ti", "cores": 4352, "vram": 8},
    {"time": 44, "gpu": "RTX3060 Ti", "cores": 4864, "vram": 8},
    {"time": 46, "gpu": "RTX4060 Ti 16GB", "cores": 4352, "vram": 16},
    {"time": 48, "gpu": "RTX3060 Ti", "cores": 4864, "vram": 8},
    {"time": 50, "gpu": "RTX4060", "cores": 3072, "vram": 8},
    {"time": 50, "gpu": "RTX2080 Super", "cores": 3072, "vram": 8},
    {"time": 52, "gpu": "RTX4060", "cores": 3072, "vram": 8},
    {"time": 54, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 54, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 54, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 54, "gpu": "RTX4060", "cores": 3072, "vram": 8},
    {"time": 56, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 56, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 56, "gpu": "RTX4070 Mobile", "cores": 4608, "vram": 8},
    {"time": 58, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 58, "gpu": "RTX2070 Super", "cores": 2560, "vram": 8},
    {"time": 58, "gpu": "RTX4060 Mobile", "cores": 3072, "vram": 8},
    {"time": 58, "gpu": "RTX3060 Mobile", "cores": 3840, "vram": 6},
    {"time": 60, "gpu": "RTX3060 12GB", "cores": 3584, "vram": 12},
    {"time": 66, "gpu": "Quadro RTX4000", "cores": 2304, "vram": 8},
    {"time": 70, "gpu": "Quadro RTX4000", "cores": 2304, "vram": 8},
    {"time": 70, "gpu": "RTX3050 8GB", "cores": 2560, "vram": 8},
    {"time": 70, "gpu": "RTX2060", "cores": 1920, "vram": 6},
]

for data, filename in [(data_all, "cuda_vs_time_all.png"), (data_filtered, "cuda_vs_time_filtered.png")]:
    df = pd.DataFrame(data)

    vram_palette = {
        4:  "#e6194b",
        6:  "#f58231",
        8:  "#3cb44b",
        10: "#4363d8",
        12: "#911eb4",
        16: "#42d4f4",
        24: "#f032e6",
    }

    z = np.polyfit(df["cores"], df["time"], 1)
    p = np.poly1d(z)
    x_line = np.linspace(df["cores"].min(), df["cores"].max(), 300)

    fig, ax = plt.subplots(figsize=(18, 11))

    plotted_vram = set()
    for _, row in df.iterrows():
        vram = int(row["vram"])
        color = vram_palette[vram]
        label = f"{vram} GB" if vram not in plotted_vram else None
        ax.scatter(row["cores"], row["time"],
                   s=90, color=color, edgecolors="white",
                   linewidths=0.5, zorder=3, label=label)
        plotted_vram.add(vram)

    ax.plot(x_line, p(x_line),
            color="#555555", linewidth=1.6, linestyle="--",
            label=f"Trend",
            zorder=2)

    for _, row in df.iterrows():
        ax.annotate(
            row["gpu"],
            xy=(row["cores"], row["time"]),
            xytext=(6, 3),
            textcoords="offset points",
            fontsize=6.5,
            color="#333333",
        )

    ax.set_title("NVIDIA GPU: CUDA Cores vs Processing Time  (color = VRAM)", fontsize=15, pad=14)
    ax.set_xlabel("CUDA Cores", fontsize=13)
    ax.set_ylabel("Processing Time (seconds)", fontsize=13)

    handles, labels = ax.get_legend_handles_labels()
    vram_pairs = sorted(
        [(h, l) for h, l in zip(handles, labels) if "GB" in l],
        key=lambda x: int(x[1].split()[0])
    )
    trend_pairs = [(h, l) for h, l in zip(handles, labels) if "Trend" in l]
    sh, sl = zip(*(vram_pairs + trend_pairs))
    ax.legend(sh, sl, title="VRAM", fontsize=10, title_fontsize=10)

    ax.grid(True, alpha=0.25)
    plt.tight_layout()
    plt.savefig(filename, dpi=300, bbox_inches="tight")