import types
import yaml
import logging
import os
import tempfile
import shutil
import datetime

"""Parameter file"""


def fetch_options(idx_dataset=0):
    """Define default values for parameters used within tf"""
    op = types.SimpleNamespace()

    # setup
    op.device = '0'  # standard device, can be bypassed if device is provided during call

    # i/o
    op.path_parser_cfg = '/home/io_patterns_nako.yaml'

    if idx_dataset == 0:
        op.dir_data = '/home/NAKO/'
    else:
        raise ValueError('Chosen dataset idx is not available')

    op.set_split = 70

    # pipeline & augmentation
    op.shape_image = [320, 260, 316]
    op.shape_image_eval = [160, 160, 160]  # smaller image results in faster eval time (less patches) old:[240, 240, 240], i have changed it to [160, 160, 160]
    op.shape_input = [64, 64, 64]  # has to correspond to your designed model
    op.shape_output = [16, 16, 16]  # has to correspond to your designed model
    op.size_batch = 24
    op.size_batch_eval = 50
    op.size_buffer = 1
    op.num_parallel_calls = 4
    op.repeat = 5
    op.b_shuffle = True
    op.b_eval_labels_patch = False
    op.b_eval_labels_image = True

    op.patches_per_class = [2, 1, 1, 1, 1, 1]  # Note: atm has to be reflected in your input pipeline #2

    # patch augmentation
    op.sigma_offset = 0.1
    op.sigma_noise = 0.05
    op.sigma_pos = 0.08

    # image augmentation
    op.b_mirror = False
    op.b_rotate = True
    op.b_scale = True
    op.b_warp = False
    op.b_permute_labels = False
    op.angle_max = 7
    op.scale_factor = 0.08
    op.delta_max = 0

    # model
    op.channels = 2
    op.channels_out = 6
    op.b_dynamic_pos_mid = True
    op.b_dynamic_pos_end = False
    op.filters = 32
    op.dense_layers = 2
    op.alpha = 0.2
    op.rate_dropout = 0.0

    # optimizer
    op.rate_learning = 0.00001  # std: 1e-2 - 1e-6, too low: slow learning
    op.beta1 = 0.9  # std: 0.9
    op.beta2 = 0.999  # std: 0.999
    op.epsilon = 0.00000001  # std: 1e-8, too high: slow learning

    # session
    op.num_epochs = 701
    op.b_continuous_eval = True
    op.b_restore = False
    op.save_summary_steps = 1
    op.b_viewer_train = False
    op.b_viewer_eval = False
    op.b_save_pred = False

    # logging
    op.log_level = logging.INFO

    # seed
    op.b_use_seed = False
    op.random_seed = 100

    return op


class Params(object):

    def __init__(self, path_yaml='', model_dir='', idx_dataset=-1, b_recreate=False):

        # set yaml path:
        self.path_yaml = path_yaml

        # fetch dataset idx to choose hardcoded split
        self.idx_dataset = idx_dataset

        # add passed/generated params file
        self.update(b_recreate)

        # add default values
        if model_dir:
            self.set_model_dir(model_dir)

    def create(self):
        # Create new default params
        logging.info('Creating a new set of parameters in %s', self.path_yaml)
        self.__dict__.update(fetch_options(idx_dataset=self.idx_dataset).__dict__)
        self.save()

    def save(self):
        if not self.path_yaml:
            _, self.path_yaml = tempfile.mkstemp()
        try:
            with open(self.path_yaml, 'w') as file:
                yaml.dump(self.__dict__, file, indent=4)
        except Exception:
            os.remove(self.path_yaml)
            raise Exception()

    def update(self, b_recreate):
        if b_recreate or not os.path.isfile(self.path_yaml):
            self.create()

        with open(self.path_yaml, 'r') as file:
            params = yaml.load(file)
            self.__dict__.update(params)

    def set_path(self, path_yaml):
        self.path_yaml = path_yaml

    def set_model_dir(self, model_dir):
        self.__dict__['date'] = datetime.datetime.now().strftime('_%Y-%m-%dT%H-%M-%S')
        self.__dict__['dir_model'] = os.path.join(model_dir, 'run' + self.__dict__['date'])
        self.__dict__['dir_logs_train'] = os.path.join(self.__dict__['dir_model'], 'logs', 'train')
        self.__dict__['dir_logs_eval'] = os.path.join(self.__dict__['dir_model'], 'logs', 'eval')
        self.__dict__['dir_graphs_train'] = os.path.join(self.__dict__['dir_model'], 'graphs', 'train')
        self.__dict__['dir_graphs_eval'] = os.path.join(self.__dict__['dir_model'], 'graphs', 'eval')
        self.__dict__['dir_ckpts'] = os.path.join(self.__dict__['dir_model'], 'ckpts')
        self.__dict__['dir_ckpts_best'] = os.path.join(self.__dict__['dir_model'], 'ckpts_best')

        # Create environment
        for k, v in self.__dict__.items():
            if 'dir' in k:
                if not os.path.exists(v):
                    os.makedirs(v, exist_ok=True)

        # Move params to corresponding model folder
        self.move(os.path.join(self.__dict__['dir_model'], 'params.yaml'))

        # Save with dirs
        self.save()

    def move(self, path_yaml_new):
        shutil.move(self.path_yaml, path_yaml_new)
        self.path_yaml = path_yaml_new

    @property
    def dict(self):
        """Gives dict-like access to Params instance by `params.dict['learning_rate']`"""
        return self.__dict__