#!/usr/bin/env python3
import csv
import re
import subprocess
import sys
from itertools import product
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
BUILD = ROOT / "kernels_op_7" / "build"
CSV_PATH = ROOT / "kernels_op_7" / "sweep.csv"

# Just BR x BC. 9 configs. BR=128 will OOM on smem (48KB cap), expected NOT_FOUND.
GRID = {
    "br": [32, 64, 128],
    "bc": [32, 64, 128],
}

# from Makefile
NVCC = "nvcc"
CFLAGS = "-O3 -arch=sm_86 -std=c++17 -rdc=true -g -lineinfo"
INCLUDES = "-I/usr/local/cuda/include -I/work/hdd/bche/hshah6/cutlass/include"
LIBS = "-L/usr/local/cuda/lib64 -lcublas -lcublasLt"

KERNEL_NAME = "flash_attention_forward_kernel"


def run(cmd):
    return subprocess.run(cmd, shell=True, cwd=ROOT, capture_output=True, text=True)


def compile_one(tag, defs):
    """Two-step compile matching the Makefile."""
    obj = BUILD / f"{tag}.o"
    binary = BUILD / tag
    r = run(f"{NVCC} {CFLAGS} {INCLUDES} {defs} -c test_gpt2.cu -o {obj}")
    if r.returncode != 0:
        return None, f"COMPILE FAILED:\n{r.stderr}"
    r = run(f"{NVCC} {CFLAGS} -o {binary} {obj} {LIBS}")
    if r.returncode != 0:
        return None, f"LINK FAILED:\n{r.stderr}"
    return binary, ""


def profile_one(binary, tag):
    report = BUILD / tag
    log = BUILD / f"{tag}.log"
    cmd = (
        f"nsys profile --stats=true --force-overwrite=true "
        f"--cuda-event-trace=false -o {report} {binary}"
    )
    r = run(cmd)
    log.write_text((r.stdout or "") + "\n" + (r.stderr or ""))
    return log


def parse_avg_ns(log_path, kernel_name):
    pattern = re.compile(rf"^\s*\d.*\b{re.escape(kernel_name)}\b")
    for line in log_path.read_text().splitlines():
        if pattern.search(line):
            cols = line.split()
            if len(cols) >= 4:
                return cols[3]
    return "NOT_FOUND"


def main():
    BUILD.mkdir(parents=True, exist_ok=True)
    configs = list(product(GRID["br"], GRID["bc"]))

    first_fail_printed = False
    with CSV_PATH.open("w", newline="") as fp:
        out = csv.writer(fp)
        out.writerow(["BR", "BC", "fa2_avg_ns"])

        for i, (br, bc) in enumerate(configs, 1):
            tag = f"br{br}_bc{bc}"

            defs = f"-DSWEEP_BR={br} -DSWEEP_BC={bc}"
            binary, err = compile_one(tag, defs)

            if binary is None:
                if not first_fail_printed:
                    print(err)
                    first_fail_printed = True
                out.writerow([br, bc, "BUILD_FAIL"])
                fp.flush()
                continue

            log = profile_one(binary, tag)
            out.writerow([br, bc, parse_avg_ns(log, KERNEL_NAME)])
            fp.flush()

    return 0


if __name__ == "__main__":
    sys.exit(main())