#!/usr/bin/env python
from __future__ import annotations

import argparse
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.experiments.baselines import (  # noqa: E402
    list_baselines,
    prepare_dataset_for_baseline,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Prepare a transformed CIL baseline dataset.")
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--baseline", choices=list_baselines(), required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--shard-size", type=int, default=2048)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args(argv)

    output = prepare_dataset_for_baseline(
        args.dataset,
        args.baseline,
        args.out,
        shard_size=args.shard_size,
        seed=args.seed,
    )
    print(output)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())