#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.12"
# dependencies = ["matplotlib", "PySide6"]
# ///

import argparse
import csv
import matplotlib
matplotlib.use("QtAgg")
import matplotlib.pyplot as plt
from pathlib import Path

def parse_args():
    parser = argparse.ArgumentParser(description="Plot three CSV columns.")
    parser.add_argument("csv_path", help="Path to the CSV file")
    parser.add_argument("--x", default=0, help="Column for the x axis")
    parser.add_argument("--y0", default=1, help="Column 0 for the y axis")
    parser.add_argument("--y1", default=2, help="Column 1 for the y axis")
    parser.add_argument("--save", action="store_true", help="Save the chart instead of displaying it")
    args = parser.parse_args()
    return args

def to_ns(datum):
    if datum.endswith(" ns"):
        return float(datum[:-3])
    if datum.endswith(" us"):
        return float(datum[:-3]) * 1000
    if datum.endswith(" ms"):
        return float(datum[:-3]) * 1000 * 1000
    # Assume it's already in ns
    return float(datum)

def get_basic_save_path(args):
    return Path(args.csv_path).with_suffix(".png")

def get_ratio_save_path(args):
    path = Path(args.csv_path)
    stem = path.stem
    return Path(args.csv_path).with_stem(stem + "_ratio").with_suffix(".png")

def parse_data(args):
    data = ([], [], [])
    with open(args.csv_path, newline="", encoding="utf-8") as file:
      reader = (row for row in csv.reader(file))
      header = next(reader)

      for row in reader:
          x = float(row[args.x])
          y0 = to_ns(row[args.y0])
          y1 = to_ns(row[args.y1])

          data[0].append(x)
          data[1].append(y0)
          data[2].append(y1)

    return data

def plot_basic(args, data):
    x, y0, y1 = data
    fig, ax = plt.subplots()
    ax.plot(x, y0, marker="o", label="rustfft")
    ax.plot(x, y1, marker="o", label="vfft")
    ax.set(xlabel="Input length", ylabel="Time (ns)")
    ax.legend()
    ax.grid(True, alpha=0.3)
    fig.tight_layout()

    if args.save:
      fig.savefig(get_basic_save_path(args), dpi=150)
    else:
      plt.show()

def plot_ratio(args, data):
    x, y0, y1 = data
    ratio = [ x1 / x0 for x0, x1 in zip(y0, y1) ]
    fig, ax = plt.subplots()
    ax.plot(x, ratio, marker="o", label="vfft / rustfft")
    ax.set_xlim(left=0)
    ax.set_ylim(bottom=0)
    ax.axhline(y=1, color="red", linestyle="--")
    ax.set(xlabel="Input length", ylabel="Execution time ratio (vfft / rustfft)")
    ax.legend()
    ax.grid(True, alpha=0.3)
    fig.tight_layout()

    if args.save:
      fig.savefig(get_ratio_save_path(args), dpi=150)
    else:
      plt.show()

def print_stats(data):
    x, y0, y1 = data

    num_wins = 0
    for x0, x1 in zip(y0, y1):
        if x1 < x0:
            num_wins += 1

    print(f"num trials: {len(x)}")
    print(f"num wins:   {num_wins}")
    print(f"win rate:   {num_wins / len(x)}")

def main():
    args = parse_args()
    data = parse_data(args)
    plot_basic(args, data)
    plot_ratio(args, data)
    print_stats(data)


if __name__ == "__main__":
    main()