import argparse
import re
from pathlib import Path
import pandas as pd

FILENAME_RE = re.compile(
    r"^log_(?P<alg>.+?)-n_workers(?P<n_workers>\d+)-(?P<comput>.+?)-dataset(?P<dataset>.+?)-partition(?P<partition>.+?)-model(?P<model>.+?)-alpha(?P<alpha>.+?)-seed(?P<seed>\d+)-(?P<metric>test-acc|train-loss)\.csv$"
)


def parse_filename(path: Path):
    match = FILENAME_RE.match(path.name)
    if not match:
        return None
    info = match.groupdict()
    info["seed"] = int(info["seed"])
    info["n_workers"] = int(info["n_workers"])
    return info


def metric_suffix(metric_name: str) -> str:
    return "Test/Acc" if metric_name == "test-acc" else "Train/Loss"


def metric_tag(metric_name: str) -> str:
    return "acc" if metric_name == "test-acc" else "loss"


def build_column_prefix(info: dict) -> str:
    return f"{info['alg']}-{info['comput']}-alpha{info['alpha']}"


def merge_seed_metric_files(seed_files, output_path: Path):
    merged = None
    for path, info in sorted(seed_files, key=lambda x: x[0].name):
        metric_col = metric_suffix(info["metric"])
        prefix = build_column_prefix(info)
        out_col = f"{prefix} - {metric_col}"

        df = pd.read_csv(path)
        if "Step/Comput" not in df.columns or metric_col not in df.columns:
            raise ValueError(
                f"File {path} must contain 'Step/Comput' and '{metric_col}' columns."
            )

        cur = df[["Step/Comput", metric_col]].dropna(axis=0, how="any").copy()
        cur = cur.sort_values("Step/Comput")
        cur = cur.drop_duplicates(subset=["Step/Comput"], keep="last")
        cur = cur.rename(columns={metric_col: out_col})

        if merged is None:
            merged = cur
        else:
            merged = pd.merge(merged, cur, on="Step/Comput", how="outer")

    if merged is None:
        raise ValueError(f"No files were provided for output {output_path}.")

    merged = merged.sort_values("Step/Comput").reset_index(drop=True)
    merged.insert(0, "Step/Comput/Cul", merged["Step/Comput"])
    output_path.parent.mkdir(parents=True, exist_ok=True)
    merged.to_csv(output_path, index=False)


def main():
    parser = argparse.ArgumentParser(
        description="Merge raw per-run CSV logs into seed-level comparison CSV files used by draw_csv.py."
    )
    parser.add_argument("--logdir", type=str, default="./results/raw", help="Directory containing raw log CSV files.")
    parser.add_argument(
        "--output-dir",
        type=str,
        default="./results/raw",
        help="Directory for merged comparison CSV files. The default matches the repository workflow and stores them in results/raw/.",
    )
    parser.add_argument("--dataset", type=str, default="cifar10", help="Dataset name to filter from raw logs.")
    parser.add_argument(
        "--dataset-tag",
        type=str,
        default=None,
        help="Dataset tag to use in the merged CSV filenames. Defaults to the dataset name with '10' removed when applicable, e.g. cifar10 -> cifar.",
    )
    parser.add_argument("--n-workers", type=int, required=True, help="Number of workers to filter.")
    parser.add_argument("--partition", type=str, default="noniid-labeldir", help="Partition name to filter.")
    parser.add_argument("--model", type=str, default="cnn", help="Model name to filter.")
    parser.add_argument(
        "--seeds",
        type=int,
        nargs="*",
        default=None,
        help="Optional list of seeds to include. If omitted, all available seeds are used.",
    )
    args = parser.parse_args()

    logdir = Path(args.logdir)
    output_dir = Path(args.output_dir)
    if not logdir.is_dir():
        raise FileNotFoundError(f"Log directory does not exist: {logdir}")

    dataset_tag = args.dataset_tag
    if dataset_tag is None:
        dataset_tag = args.dataset[:-2] if args.dataset.endswith("10") else args.dataset

    grouped = {}
    for path in sorted(logdir.glob("log_*.csv")):
        info = parse_filename(path)
        if info is None:
            continue
        if info["dataset"] != args.dataset:
            continue
        if info["n_workers"] != args.n_workers:
            continue
        if info["partition"] != args.partition:
            continue
        if info["model"] != args.model:
            continue
        if args.seeds is not None and info["seed"] not in args.seeds:
            continue
        grouped.setdefault((info["seed"], info["metric"]), []).append((path, info))

    if not grouped:
        raise RuntimeError(
            "No matching raw log CSV files were found. Please check the filters and the log directory."
        )

    for (seed, metric_name), metric_files in sorted(grouped.items()):
        out_name = f"{dataset_tag}-{metric_tag(metric_name)}-{args.n_workers}-{seed}.csv"
        output_path = output_dir / out_name
        merge_seed_metric_files(metric_files, output_path)
        print(f"Wrote {output_path}")


if __name__ == "__main__":
    main()
