#!/usr/bin/env python3
"""Create the cross validation split files and their nested training fractions."""

import argparse
import os

from auto_detect_breast_mri.config import resolve_path
from auto_detect_breast_mri.data.splits import (create_fraction_train_files,
                                               stratified_split_train_test_eval)


def main():
    """
    Create the cross validation split files and their nested training fractions.

    Paths come from the command line, or from the site config (see config.example.yaml) when an
    argument is omitted:

        python scripts/data_utils/make_splits.py --split-root <folder> --metadata-file <metadata.csv>
    """
    parser = argparse.ArgumentParser(description="Create cross validation split files.")
    parser.add_argument("--split-root", default=None,
                        help="folder the split files are written to. Config key: split_root")
    parser.add_argument("--metadata-file", default=None,
                        help="metadata export with the label columns. Config key: metadata_file")
    parser.add_argument("--fold", type=int, default=0)
    parser.add_argument("--subfold", type=int, default=0)
    parser.add_argument("--fractions", type=float, nargs="+", default=[0.05, 0.1, 0.25, 0.5],
                        help="training fractions to generate. Default: 0.05 0.1 0.25 0.5")
    args = parser.parse_args()

    split_root = resolve_path(args.split_root, "split_root", "folder for the split files")
    metadata_file = resolve_path(args.metadata_file, "metadata_file", "metadata export")
    if not split_root.endswith(os.sep):
        split_root += os.sep

    training_set_location = stratified_split_train_test_eval(split_root, metadata_file)
    training_file = training_set_location.format(args.fold, args.fold, args.subfold)
    create_fraction_train_files(training_file, args.fractions)
    print(f"Wrote split files and fractions {args.fractions} under {split_root}")


if __name__ == '__main__':
    main()