"""
Validate an already trained cropping UNet on a test split, with the training augmentation applied.

Same purpose as val_trained_unet.py in the repository root, which validates without augmentation.
Paths not given on the command line come from the site config (see config.example.yaml):

    python scripts/validation/validate_crop_unet.py --unet-path <unet_best.pth>
"""

import argparse

import torchio as tio

from auto_detect_breast_mri.config import resolve_path
from auto_detect_breast_mri.data.transforms import ImageToTensor, RandomCropOrPad, ZNormalization, masking_method
from auto_detect_breast_mri.models.unet_crop import UNet_Crop
from auto_detect_breast_mri.data.loaders import get_subjects_dataloader
from auto_detect_breast_mri.data.metadata import get_uka_metatensor
def build_parser():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--unet-path", required=True,
                        help="pth file of the trained cropping UNet")
    parser.add_argument("--data-path", default=None,
                        help="image root. Config key: data_root")
    parser.add_argument("--feature-path", default=None,
                        help="metadata export. Config key: metadata_file")
    parser.add_argument("--split-files-folder", default=None,
                        help="folder holding the split files. Config key: split_root")
    parser.add_argument("-b", "--batch-size", type=int, default=4)
    parser.add_argument("-c", "--fold", type=int, default=0, help="outer fold whose test set is used")
    parser.add_argument("-n", "--number-of-epochs", type=int, default=10,
                        help="how often the validation set is passed through. Default: 10")
    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
    return parser


if __name__ == '__main__':
    args = build_parser().parse_args()

    path_base = resolve_path(args.data_path, "data_root", "image root")
    feature_path = resolve_path(args.feature_path, "metadata_file", "metadata export")
    split_root = resolve_path(args.split_files_folder, "split_root", "folder with the split files")
    test_split = f"{split_root.rstrip('/')}/fold{args.fold}/stratified_test_set"

    protocol = ['Dyn_0']
    pre_image_shape = (32, 512, 512)
    augmentation = tio.Compose([
        tio.RandomFlip((0, 1, 2), flip_probability=0.5),
        RandomCropOrPad(pre_image_shape),
        ZNormalization(per_channel=True, percentiles=(0.5, 99.5), masking_method=masking_method),
        tio.RandomNoise(std=(0.25, 0.5)),
        ImageToTensor()
    ])

    # ------------ Initialize Model ------------
    in_channel = len(protocol)
    CropNet = UNet_Crop(path_to_unet=args.unet_path, input_channel=in_channel, running_device=args.device,
                        show_validation=True, criterion='dice')

    # Dataset and Loader
    feature_dataframe = get_uka_metatensor(0, feature_path)
    test_loader = get_subjects_dataloader(path_base, feature_dataframe, pre_image_shape,
                                          transform=augmentation, protocol=protocol,
                                          data_selection_file_sceleton=test_split,
                                          batch_size=args.batch_size, fold=args.fold,
                                          segmentation_task=True)
    CropNet.validate(args.number_of_epochs, test_loader)