"""Export the fixed-seed random split used for the main experiment."""

from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

import pandas as pd


PACKAGE_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PACKAGE_ROOT))

from concrete_gnn import random_split_indices, split_indices_sha256  # noqa: E402


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data",
        default=PACKAGE_ROOT / "data" / "HiConGNN_processed_database.csv",
        type=Path,
    )
    parser.add_argument(
        "--output",
        default=PACKAGE_ROOT / "data" / "split_random_seed17.csv",
        type=Path,
    )
    parser.add_argument("--seed", type=int, default=17)
    args = parser.parse_args()

    n_rows = len(pd.read_csv(args.data, usecols=[0]))
    train, validation, test = random_split_indices(n_rows, args.seed)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.writer(stream, lineterminator="\n")
        writer.writerow(("row_index", "split"))
        for name, indices in (
            ("train", train),
            ("validation", validation),
            ("test", test),
        ):
            writer.writerows((index, name) for index in indices)

    digest = split_indices_sha256(train, validation, test)
    print(
        f"wrote {args.output} rows={n_rows} train={len(train)} "
        f"validation={len(validation)} test={len(test)} sha256={digest}"
    )


if __name__ == "__main__":
    main()
