import argparse

import shutil

import pickle

import copy

import os

import numpy as np

class BaseArgParser(object):

    def __init__(self):

        self.parser = argparse.ArgumentParser()

    def namespace_to_dict(self, args):

        """Turns a nested Namespace object to a nested dictionary"""

        args_dict = vars(copy.deepcopy(args))

        for arg in args_dict:

            obj = args_dict[arg]

            if isinstance(obj, argparse.Namespace):

                args_dict[arg] = self.namespace_to_dict(obj)

        return args_dict

    def fix_nested_namespaces(self, args):

        """Makes sure that nested Namespace work. Supports only one level of nesting."""

        group_name_keys = []

        for key in args.__dict__:

            if '.' in key:

                group, name = key.split('.')

                group_name_keys.append((group, name, key))

        for group, name, key in group_name_keys:

            if group not in args:

                args.__dict__[group] = argparse.Namespace()

            args.__dict__[group].__dict__[name] = args.__dict__[key]

            del args.__dict__[key]

    def parse_args(self):

        args = self.parser.parse_args()

        args = self.namespace_to_dict(args)

        self.fix_nested_namespaces(args)

        return args

class PreproArgParser(BaseArgParser):

    def __init__(self):

        super(PreproArgParser, self).__init__()

        self.parser.add_argument('--in_locs', type=str, required=True,

                help='Comma-separated list of paths to all data folders.')

        self.parser.add_argument('--modalities', type=str, required=True,

                help='Comma-separated list of all input modalities to use.')

        self.parser.add_argument('--truth', type=str, required=True,

                help='Truth label pattern to use.')

        self.parser.add_argument('--create_val', action='store_true', default=False,

                help='Whether to create validation set.')

        self.parser.add_argument('--out_loc', type=str, default='./data',

                help='Location to write preprocessed data.')

    def parse_args(self):

        args = self.parser.parse_args()

        # Create list of all input datasets.

        args.in_locs = args.in_locs.split(',')

        # Create list of all accepted modalities.

        args.modalities = args.modalities.split(',')

        # Create output directory if it doesn't already exist.

        if os.path.isdir(args.out_loc):

            shutil.rmtree(args.out_loc)

        args.train_loc = os.path.join(args.out_loc, 'train')

        args.val_loc = os.path.join(args.out_loc, 'val')

        os.mkdir(args.out_loc)

        os.mkdir(args.train_loc)

        os.mkdir(args.val_loc)

        return args

class TrainArgParser(BaseArgParser):

    def __init__(self):

        super(TrainArgParser, self).__init__()

        # Data args.

        self.parser.add_argument('--train_loc', type=str, required=True,

                help='Location of .tfrecords training data.')

        self.parser.add_argument('--prepro_loc', type=str, required=True,

                help='Location of preprocessed dump.')

        self.parser.add_argument('--val_loc', type=str, default='',

                help='Location of .tfrecords validation data.')

        # Checkpoint args.

        self.parser.add_argument('--save_folder', type=str, default='',

                help='Output folder to save checkpoints, logs, and configs.')

        self.parser.add_argument('--load_folder', type=str, default='',

                help='Input folder to load checkpoints and configs to resume training.')

        # Training args.

        self.parser.add_argument('--lr', type=float, default=1e-4,

                help='Initial learning rate for training.')

        self.parser.add_argument('--batch_size', type=int, default=1,

                help='Batch size to use in training.')

        self.parser.add_argument('--patience', type=int, default=-1,

                help='Number of epochs without validation improvement to stop training.')

        self.parser.add_argument('--n_epochs', type=int, default=300,

                help='Number of epochs to train for.')

        self.parser.add_argument('--gpu', action='store_true', default=False,

                help='Whether to train using GPU.')

        # Augmentation args.

        self.parser.add_argument('--crop_size', type=str, default='128,128,128',

                help='Crop size of image (comma-separated h,w,d).')

        # Model args.

        self.parser.add_argument('--data_format', type=str, dest='model_args.data_format',

                default='channels_first', choices=['channels_last', 'channels_first'],

                help='Data format to be passed through the model.')

        self.parser.add_argument('--base_filters', type=int, dest='model_args.base_filters', default=32,

                help='Number of filters in the base convolutional layer.')

        self.parser.add_argument('--depth', type=int, dest='model_args.depth', default=4,

                help='Number of spatial levels through the model.')

        self.parser.add_argument('--l2_scale', type=float, dest='model_args.l2_scale', default=1e-5,

                help='Scale of L2 regularization applied to all kernels.')

        self.parser.add_argument('--dropout', type=float, dest='model_args.dropout', default=0.2,

                help='Dropout ratio to apply to input data.')

        self.parser.add_argument('--groups', type=int, dest='model_args.groups', default=8,

                help='Number of groups in group normalization.')

        self.parser.add_argument('--reduction', type=int, dest='model_args.reduction', default=8,

                help='Size of reduction ratio in squeeze-excitation layers.')

        self.parser.add_argument('--downsampling', type=str, dest='model_args.downsampling',

                default='conv', choices=['conv', 'max', 'avg'],

                help='Type of downsampling method.')

        self.parser.add_argument('--upsampling', type=str, dest='model_args.upsampling',

                default='conv', choices=['conv', 'linear'],

                help='Type of upsampling method.')

        self.parser.add_argument('--out_ch', type=int, dest='model_args.out_ch', default=3,

                help='Number of output classes.')

    def parse_args(self):

        args = self.parser.parse_args()

        # Fix nested Namespaces.

        self.fix_nested_namespaces(args)

        args.data_format = args.model_args.data_format

        # Check data format and GPU compatibility.

        args.device = '/device:GPU:0' if args.gpu else '/cpu:0'

        if not args.gpu:

            assert args.model_args.data_format == 'channels_last', \

                'tf.keras.layers.Conv3D only supports `channels_last` input for CPU.'

        # Convert model args to dictionaries.

        args.model_args = self.namespace_to_dict(args.model_args)

        # Set crop size.

        args.crop_size = args.crop_size.split(',')

        args.crop_size = [int(s) for s in args.crop_size]

        # Load preprocessed stats.

        prepro = np.load(args.prepro_loc).item()

        args.prepro_size = [prepro['size']['h'], prepro['size']['w'], prepro['size']['d'], prepro['size']['c']]

        # Check that sizes work out.

        assert (args.model_args['base_filters'] / 2) % args.model_args['groups'] == 0, \

            'Base filters must be a multiple of {} for group normalization at lowest spatial level.'.format(args.model_args['groups'] * 2)

        assert args.model_args['base_filters'] % args.model_args['reduction'] == 0, \

            'Base filters must be a multiple of {} for squeeze-excitation reduction.'.format(args.model_args['reduction'])

        # Add args.model_args.in_ch for output size of variational autoencoder.

        args.model_args['in_ch'] = prepro['size']['c']

        # Check for checkpointing option.

        if args.load_folder:

            with open(os.path.join(args.load_folder, 'train_args.pkl'), 'rb') as f:

                chkpt_args = pickle.load(f)

                args.model_args = chkpt_args['model_args']

                args.crop_size = chkpt_args.crop_size

                assert isinstance(args.model_args, dict)

            args.save_folder = args.load_folder

        # Create checkpoint folder if necessary.

        if not os.path.isdir(args.save_folder):

            os.mkdir(args.save_folder)

        # Save training args.

        with open(os.path.join(args.save_folder, 'train_args.pkl'), 'wb') as f:

            pickle.dump(self.namespace_to_dict(args), f)

        return args

class TestArgParser(BaseArgParser):

    def __init__(self):

        super(TestArgParser, self).__init__()

        # Data.

        self.parser.add_argument('--in_locs', type=str, required=True,

                help='Comma-separated paths of test data.')

        self.parser.add_argument('--modalities', type=str, required=True,

                help='Comma-separated modalities to be used as input')

        self.parser.add_argument('--truth', type=str, default='',

                help='Truth label pattern to use (optional).')

        # Training and preprocessing stats.

        self.parser.add_argument('--tumor_prepro', type=str, required=True,

                help='Path to Numpy preprocessing dump for tumor segmentation.')

        self.parser.add_argument('--skull_prepro', type=str, default='',

                help='Path to Numpy preprocessing dump for skull segmentation.')

        self.parser.add_argument('--tumor_model', type=str, required=True,

                help='Path to checkpoint folder for tumor segmentation.')

        self.parser.add_argument('--skull_model', type=str, default='',

                help='Path to checkpoint folder for skull-stripping segmentation.')

        # Input normalization parameters.

        self.parser.add_argument('--order', type=int, default=3,

                help='Order of interpolation function to be used in voxel resizing.')

        self.parser.add_argument('--mode', type=str, default='reflect',

                help='Method of handling image edges in interpolation.')

        # Test time augmentation and segmentation.

        self.parser.add_argument('--spatial_tta', action='store_true', default=True,

                help='Whether to apply spatial augmentation on all spatial axes.')

        self.parser.add_argument('--channel_tta', type=int, default=0,

                help='Additional intensity shifting samples to take.')

        self.parser.add_argument('--threshold', type=float, default=0.5,

                help='Threshold at which to create mask from probabilities.')

        self.parser.add_argument('--gpu', action='store_true', default=False,

                help='Whether to evaluate on GPU.')

    def parse_args(self):

        args = self.parser.parse_args()

        args.modalities = args.modalities.split(',')

        args.in_locs = args.in_locs.split(',')

        # Assert proper combination of inputs.

        assert args.threshold > 0 and args.threshold < 1, \

            'Threshold must be a probability between (0, 1).'

        if args.skull_model:

            assert args.skull_prepro, 'Need skull preprocessing stats if model is provided.'

        args.skull_strip = bool(args.skull_model)

        # Load model args.

        with open(os.path.join(args.tumor_model, 'train_args.pkl'), 'rb') as f:

            train_args = pickle.load(f)

            args.tumor_model_args = train_args['model_args']

            args.tumor_spatial_res = 2 ** args.tumor_model_args['depth']

            args.tumor_crop_size = train_args['crop_size']

        if args.skull_model:

            with open(os.path.join(args.skull_model, 'train_args.pkl'), 'rb') as f:

                train_args = pickle.load(f)

                args.skull_model_args = train_args['model_args']

                args.skull_spatial_res = 2 ** args.skull_model_args['depth']

                args.skull_crop_size = train_args['crop_size']

        # Load prepro stats.

        args.tumor_prepro = np.load(args.tumor_prepro).item()

        if args.skull_prepro:

            args.skull_prepro = np.load(args.skull_prepro).item()

        # Check data format and GPU compatibility.

        args.device = '/device:GPU:0' if args.gpu else '/cpu:0'

        if not args.gpu:

            assert args.tumor_model_args['data_format'] == 'channels_last' and args.skull_model_args['data_format'], \

                'tf.keras.layers.Conv3D only supports `channels_last` input for CPU.'

        return args