import json

import os

from glob import glob

from pathlib import Path

import numpy as np

import nibabel as nib

from brats.preprocess import window_intensities_data

from fetal_net.utils.utils import get_image

from fetal_net.postprocess import postprocess_prediction as process_pred

from brats.utils import get_last_model_path

from fetal_net.normalize import normalize_data

from fetal_net.prediction import run_validation_cases, patch_wise_prediction

import argparse

from fetal_net.training import load_old_model

from fetal_net.utils.cut_relevant_areas import find_bounding_box, cut_bounding_box, check_bounding_box

original_data_folder = '../Datasets/Fetus'

def main(pred_dir, config, split='test', overlap_factor=1, preprocess_method=None):

    padding = [16, 16, 8]

    prediction2_dir = os.path.abspath(os.path.join(config['base_dir'], 'predictions2', split))

    model = load_old_model(get_last_model_path(config["model_file"]))

    with open(os.path.join(opts.config_dir, 'norm_params.json'), 'r') as f:

        norm_params = json.load(f)

    for sample_folder in glob(os.path.join(pred_dir, split, '*')):

        mask_path = os.path.join(sample_folder, 'prediction.nii.gz')

        truth_path = os.path.join(sample_folder, 'truth.nii.gz')

        subject_id = Path(sample_folder).name

        dest_folder = os.path.join(prediction2_dir, subject_id)

        Path(dest_folder).mkdir(parents=True, exist_ok=True)

        truth = nib.load(truth_path)

        nib.save(truth, os.path.join(dest_folder, Path(truth_path).name))

        mask = nib.load(mask_path)

        mask = process_pred(mask.get_data(), gaussian_std=0.5, threshold=0.5)

        bbox_start, bbox_end = find_bounding_box(mask)

        check_bounding_box(mask, bbox_start, bbox_end)

        if padding is not None:

            bbox_start = np.maximum(bbox_start - padding, 0)

            bbox_end = np.minimum(bbox_end + padding, mask.shape)

        print("BBox: {}-{}".format(bbox_start, bbox_end))

        volume = nib.load(os.path.join(original_data_folder, subject_id, 'volume.nii'))

        orig_volume_shape = np.array(volume.get_data().shape)

        volume = cut_bounding_box(volume, bbox_start, bbox_end).get_data().astype(np.float)

        if preprocess_method is not None:

            print('Applying preprocess by {}...'.format(preprocess_method))

            if preprocess_method == 'window_1_99':

                volume = window_intensities_data(volume)

            else:

                raise Exception('Unknown preprocess: {}'.format(preprocess_method))

        if norm_params is not None and any(norm_params.values()):

            volume = normalize_data(volume, mean=norm_params['mean'], std=norm_params['std'])

        prediction = patch_wise_prediction(

            model=model, data=np.expand_dims(volume, 0),

            patch_shape=config["patch_shape"] + [config["patch_depth"]],

            overlap_factor=overlap_factor

        )

        prediction = prediction.squeeze()

        padding2 = list(zip(bbox_start, orig_volume_shape - bbox_end))

        print(padding2)

        prediction = np.pad(prediction, padding2, mode='constant', constant_values=0)

        assert all([s1 == s2 for s1, s2 in zip(prediction.shape, orig_volume_shape)])

        prediction = get_image(prediction)

        nib.save(prediction, os.path.join(dest_folder, Path(mask_path).name))

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument("--pred_dir", help="specifies config dir path",

                        type=str, required=True)

    parser.add_argument("--config_dir", help="specifies config dir path",

                        type=str, required=True)

    parser.add_argument("--preprocess", help="specifies config dir path",

                        type=str, required=True)

    parser.add_argument("--split", help="What split to predict on? (test/val)",

                        type=str, default='test')

    parser.add_argument("--overlap_factor", help="specifies overlap between prediction patches",

                        type=float, default=1)

    opts = parser.parse_args()

    with open(os.path.join(opts.config_dir, 'config.json')) as f:

        config = json.load(f)

    main(opts.pred_dir, config, opts.split, opts.overlap_factor, opts.preprocess)