import os

import torchio as tio

from auto_detect_breast_mri.config import resolve_path
from auto_detect_breast_mri.data.breast_mri_dataset import BreastMRISubjects
#####################################################
# Check proportion of cases in all provided subsets #
# Subsets are specified by patientID and side in csv#
# TODO: change path_base and feature path accordingly
#####################################################


path_base = resolve_path(None, "data_root", "root folder of the NIfTI data")
pre_image_shape = (32, 512, 512)
batch_size = 1
protocol = ["Sub_1"]

parent_folder = path_base.replace(path_base.split(os.sep)[-1], "")
#subset_files = ["all.csv", "training_set.csv", "test_set.csv", "evaluation_set.csv"]
subset_files = ["small_subset.csv"]
for file in subset_files:
    traindata_set = BreastMRISubjects(path_base, parent_folder + file, protocol=protocol)
    data_loader = tio.SubjectsLoader(traindata_set, batch_size=batch_size, shuffle=True)
    false_count = 0
    true_count = 0
    for batch_id, batch in enumerate(data_loader):
        if batch['label'][0]:
            true_count += 1
        else:
            false_count += 1
    print('-------------------------------------')
    print('FILE: ', file)
    print('Number of non cancer cases: ', false_count)
    print('Number of cancer cases: ', true_count)
    print('-------------------------------------')