#!/usr/bin/env python3
"""
Verify the nesting guarantee of the training fraction files: every smaller fraction must be a
subset of every larger one, for every fold.

    python tools/check_split_fractions.py [--split-root <folder>] [--folds 5]

The split root comes from the site config when it is not given (see config.example.yaml).
"""

import argparse
import sys
from pathlib import Path

import pandas as pd

from auto_detect_breast_mri.config import column, resolve_path


def fraction_of(filename: str) -> float:
    """0.05 from '..._frac0.05.csv'."""
    return float(filename.split("_frac")[-1].rsplit(".csv", 1)[0])


def check_fold(fold_dir: Path, id_column: str, expected_files: int) -> bool:
    """True if every fraction file in fold_dir is nested in all larger ones."""
    fraction_files = sorted(fold_dir.glob("*_frac0.*.csv"), key=lambda path: fraction_of(path.name))
    if len(fraction_files) != expected_files:
        print(f"WARNING: {fold_dir.name}: expected {expected_files} fraction files, "
              f"found {len(fraction_files)}: {[path.name for path in fraction_files]}", file=sys.stderr)
    ids = {fraction_of(path.name): set(pd.read_csv(path)[id_column]) for path in fraction_files}

    ok = True
    for child, child_ids in sorted(ids.items()):
        for parent, parent_ids in sorted(ids.items()):
            if parent <= child:
                # a larger set can not be expected inside a smaller one; also skips parent == child
                continue
            missing = child_ids - parent_ids
            if missing:
                ok = False
                print(f"ERROR: {fold_dir.name}: fraction {child} is NOT contained in {parent}: "
                      f"{len(missing)} entries missing from the parent.", file=sys.stderr)
            else:
                print(f"  {fold_dir.name}: fraction {child} is fully covered by {parent}.")
    return ok


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--split-root", default=None,
                        help="folder holding fold<k>/. Config key: split_root")
    parser.add_argument("--folds", type=int, default=5, help="number of folds to check. Default: 5")
    parser.add_argument("--expected-files", type=int, default=4,
                        help="fraction files expected per fold. Default: 4")
    args = parser.parse_args()

    split_root = Path(resolve_path(args.split_root, "split_root", "folder with the split files"))
    if not split_root.exists():
        sys.exit(f"ERROR: directory {split_root} does not exist.")

    id_column = column("examination_id")
    all_ok = True
    for fold in range(args.folds):
        fold_dir = split_root / f"fold{fold}"
        if not fold_dir.exists():
            print(f"WARNING: {fold_dir} does not exist, skipping.", file=sys.stderr)
            continue
        print(f"Investigating {fold_dir}")
        all_ok &= check_fold(fold_dir, id_column, args.expected_files)

    print("OK: all fractions are nested." if all_ok else "FAIL: nesting is violated, see above.")
    sys.exit(0 if all_ok else 1)


if __name__ == "__main__":
    main()