import argparse
import json
import os

import numpy as np

from brats.utils import get_last_model_path
from fetal_net.normalize import normalize_data
from fetal_net.postprocess import postprocess_prediction as process_pred
from fetal_net.prediction import patch_wise_prediction
from fetal_net.training import load_old_model
from brats.preprocess import window_intensities_data

from scipy.io import loadmat, savemat


def main(input_mat_path, output_mat_path, config, overlap_factor, model_path, preprocess_method=None, norm_params=None):
    print(model_path)
    model = load_old_model(get_last_model_path(model_path))
    print('Loading mat from {}...'.format(input_mat_path))
    mat = loadmat(input_mat_path)
    print('Predicting mask...')
    data = mat['volume'].astype(np.float)

    if preprocess_method is not None:
        print('Applying preprocess by {}...'.format(preprocess_method))
        if preprocess_method == 'window_1_99':
            data = window_intensities_data(data)
        else:
            raise Exception('Unknown preprocess: {}'.format(preprocess_method))

    if norm_params is not None and any(norm_params.values()):
        data = normalize_data(data, mean=norm_params['mean'], std=norm_params['std'])

    prediction = \
        patch_wise_prediction(model=model,
                              data=np.expand_dims(data, 0),
                              overlap_factor=overlap_factor,
                              patch_shape=config["patch_shape"] + [config["patch_depth"]])

    print('Post-processing mask...')
    if prediction.shape[-1] > 1:
        prediction = prediction[..., 1]
    prediction = prediction.squeeze()
    mat['masks'][0, 9] = \
        process_pred(prediction, gaussian_std=0, threshold=0.2)  # .astype(np.uint8)
    mat['masks'][0, 8] = \
        process_pred(prediction, gaussian_std=1, threshold=0.5)  # .astype(np.uint8)
    mat['masks'][0, 7] = \
        process_pred(prediction, gaussian_std=0.5, threshold=0.5)  # .astype(np.uint8)
    print('Saving mat to {}'.format(output_mat_path))
    savemat(output_mat_path, mat)
    print('Finished.')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--config_dir", help="specifies config dir path",
                        type=str, required=True)
    parser.add_argument("--input_mat", help="specifies mat file dir path",
                        type=str, required=True)
    parser.add_argument("--output_mat", help="specifies mat file dir path",
                        type=str, required=True)
    parser.add_argument("--overlap_factor", help="specifies overlap between prediction patches",
                        type=float, default=0.9)
    parser.add_argument("--preprocess", help="what preprocess to do",
                        type=str, default=None)
    opts = parser.parse_args()

    with open(os.path.join(opts.config_dir, 'config.json'), 'r') as f:
        _config = json.load(f)
    with open(os.path.join(opts.config_dir, 'norm_params.json'), 'r') as f:
        _norm_params = json.load(f)

    _model_path = os.path.join(opts.config_dir, os.path.basename(_config['model_file']))
    main(opts.input_mat, opts.output_mat, _config, model_path=_model_path,
         preprocess_method=opts.preprocess, norm_params=_norm_params, overlap_factor=opts.overlap_factor)