"""
File: pylinex/model/TrainingSetCreator.py
Author: Keith Tauscher
Date: 20 Jun 2018

Description: File containing class which creates sets of curves given models
             and prior distributions.
"""
from __future__ import division
import os, time, h5py
import numpy as np
from ..util import bool_types, int_types, create_hdf5_dataset, get_hdf5_value
from distpy import DistributionSet
from .Model import Model
from .LoadModel import load_model_from_hdf5_group
try:
    # this runs with no issues in python 2 but raises error in python 3
    basestring
except:
    # this try/except allows for python 2/3 compatible string type checking
    basestring = str

class TrainingSetCreator(object):
    """
    Class which creates sets of curves given models and prior distributions.
    """
    def __init__(self, model, prior_set, num_curves, file_name, seed=None,\
        allow_errors=True, verbose=True):
        """
        Creates a training set of curves from the given models and the given
        prior region.
        
        model: Model object which will be applied to sample parameters
        prior_set: DistributionSet object from which parameters will be sampled
        num_curves: integer number of curves this object will generate
        file_name: either a file location where an hdf5 file can be saved or a
                   location where it has already begun being saved
        seed: either None if no seed or a 32-bit unsigned integer
        verbose: boolean determining if message is printed after each
                 convolution (i.e. pair of beam+maps)
        """
        self.verbose = verbose
        self.allow_errors = allow_errors
        self.seed = seed
        self.file_name = file_name
        self.num_curves = num_curves
        self.model = model
        self.prior_set = prior_set
    
    @property
    def allow_errors(self):
        """
        Property storing a boolean which determines whether errors are allowed
        (True) or disallowed (False). If they are disallowed, they are raised
        directly.
        """
        if not hasattr(self, '_allow_errors'):
            raise AttributeError("allow_errors was referenced before it " +\
                "was set.")
        return self._allow_errors
    
    @allow_errors.setter
    def allow_errors(self, value):
        """
        Setter for the allow_errors switch.
        
        value: True or False
        """
        if type(value) in bool_types:
            self._allow_errors = value
        else:
            raise TypeError("allow_errors was set to a non-bool.")
    
    @property
    def seed(self):
        """
        Property storing the seed for the random number generator (for
        reproducibility).
        """
        if not hasattr(self, '_seed'):
            raise AttributeError("seed was referenced before it was set.")
        return self._seed
    
    @seed.setter
    def seed(self, value):
        """
        Setter for the seed for the random number generator (for
        reproducibility).
        
        value: either None for no seed or a 32-bit unsigned integer
        """
        if (type(value) is type(None)) or (type(value) in int_types):
            self._seed = value
        else:
            raise TypeError("seed was set to neither None nor an integer.")
    
    @property
    def model(self):
        """
        Property storing the Model object which will be used to create the
        training set curves.
        """
        if not hasattr(self, '_model'):
            raise AttributeError("model was referenced before it was set.")
        return self._model
    
    @model.setter
    def model(self, value):
        """
        Setter for the model which will be used to create the training set
        curves.
        
        value: a Model object
        """
        if isinstance(value, Model):
            self._model = value
        else:
            raise TypeError("model was set to a non-Model object.")
    
    @property
    def prior_set(self):
        """
        Property storing the DistributionSet which is used to draw model
        parameter vector realizations.
        """
        if not hasattr(self, '_prior_set'):
            raise AttributeError("prior_set was referenced before it was set.")
        return self._prior_set
    
    @prior_set.setter
    def prior_set(self, value):
        """
        Setter for the DistributionSet which is used to draw model parameter
        vector realizations.
        
        value: DistributionSet object describing the same set of parameters as
               the model which will be used to generate training set curves
        """
        if isinstance(value, DistributionSet):
            if set(value.params) == set(self.model.parameters):
                self._prior_set = value
            else:
                not_needed_parameters =\
                    set(value.params) - set(self.model.parameters)
                missing_parameters =\
                    set(self.model.parameters) - set(value.params)
                if not_needed_parameters and missing_parameters:
                    raise ValueError(("There were some parameters ({0}) " +\
                        "which were described in the prior_set but were " +\
                        "not parameters of the model and there were some " +\
                        "parameters ({1}) of the model which were not " +\
                        "described in the prior_set.").format(\
                        not_needed_parameters, missing_parameters))
                elif not_needed_parameters:
                    raise ValueError(("Some parameters ({}) which were " +\
                        "described in the prior_set were not parameters of " +\
                        "the model.").format(not_needed_parameters))
                else:
                    raise ValueError(("Some parameters ({}) which were " +\
                        "included in the model were not described in the " +\
                        "prior_set.").format(missing_parameters))
        else:
            raise TypeError("prior_set was set to a non-DistributionSet " +\
                "object.")
    
    @property
    def verbose(self):
        """
        Property storing a boolean which determines whether or not a message is
        printed after each curve is created.
        """
        if not hasattr(self, '_verbose'):
            raise AttributeError("verbose referenced before it was set.")
        return self._verbose
    
    @verbose.setter
    def verbose(self, value):
        """
        Setter for the boolean which determines whether or not a message is
        printed after each curve is created.
        
        value: True or False
        """
        if type(value) in bool_types:
            self._verbose = value
        else:
            raise TypeError("verbose was set to a non-bool.")
    
    @property
    def file_name(self):
        """
        Property storing the name of the hdf5 file in which to save the data
        generated by this object.
        """
        if not hasattr(self, '_file_name'):
            raise AttributeError("file_name referenced before it was set.")
        return self._file_name
    
    @file_name.setter
    def file_name(self, value):
        """
        Setter for the name of the hdf5 file in which to save the data
        generated by this object.
        
        value: string name of hdf5 file, which may or may not exist
        """
        if isinstance(value, basestring):
            self._file_name = value
        else:
            raise TypeError("file_name was set to a non-string.")
    
    @property
    def num_curves(self):
        """
        Property storing the integer number of curves this TrainingSetCreator
        will create.
        """
        if not hasattr(self, '_num_curves'):
            raise AttributeError("num_curves was referenced before it was " +\
                "set.")
        return self._num_curves
    
    @num_curves.setter
    def num_curves(self, value):
        """
        Setter for the integer number of curves this TrainingSetCreator will
        create.
        
        value: positive integer number of curves to generate
        """
        if type(value) in int_types:
            if value > 0:
                self._num_curves = value
            else:
                raise ValueError("num_curves was set to a non-positive " +\
                    "integer.")
        else:
            raise TypeError("num_curves was set to a non-integer.")
    
    @property
    def file(self):
        """
        Property storing the h5py File object in which the data generated by
        this object will be saved.
        """
        if not hasattr(self, '_file'):
            if os.path.exists(self.file_name):
                self._file = h5py.File(self.file_name, 'r+')
            else:
                self._file = h5py.File(self.file_name, 'w')
                self.model.fill_hdf5_group(self._file.create_group('model'))
                self.prior_set.fill_hdf5_group(\
                    self._file.create_group('prior_set'))
                self._file.create_group('parameters')
                self._file.create_group('curves')
                self._file.attrs['next_index'] = 0
                if type(self.seed) is not type(None):
                    self._file.attrs['seed'] = self.seed
        return self._file
    
    def generate(self):
        """
        Generates (or continues generating) curves for a set of training
        curves. These curves are saved in the hdf5 file at file_name as they
        are generated.
        """
        completed = self.file.attrs['next_index']
        try:
            np.random.seed(self.seed)
            for icurve in range(self.num_curves):
                curve_string = 'curve_{:d}'.format(icurve)
                parameter_draw = self.prior_set.draw()
                if icurve < completed:
                    continue
                parameter_draw = np.array([parameter_draw[parameter]\
                    for parameter in self.model.parameters])
                try:
                    curve = self.model(parameter_draw)
                except KeyboardInterrupt:
                    raise
                except:
                    if self.allow_errors:
                        curve = np.array([np.nan])
                    else:
                        raise
                create_hdf5_dataset(self.file['parameters'], curve_string,\
                    data=parameter_draw)
                create_hdf5_dataset(self.file['curves'], curve_string,\
                    data=curve)
                completed += 1
                self.file.attrs['next_index'] = completed
                self.close()
                if self.verbose:
                    print("Finished curve #{0:d}/{1:d} at {2!s}.".format(\
                        completed, self.num_curves, time.ctime()))
                low_floor =\
                    int(np.floor(((100 * completed) - 50) / self.num_curves))
                high_floor =\
                    int(np.floor(((100 * completed) + 50) / self.num_curves))
                if (low_floor != high_floor) and\
                    (completed != self.num_curves):
                    print("Done with {0:d}% of {1:d} curves at {2!s}.".format(\
                        high_floor, self.num_curves, time.ctime()))
        except KeyboardInterrupt:
            completed_according_to_file = self.file.attrs['next_index']
            curve_string_to_delete_if_present =\
                'curve_{:d}'.format(completed_according_to_file)
            parameters_string =\
                'parameters/{!s}'.format(curve_string_to_delete_if_present)
            curves_string =\
                'curves/{!s}'.format(curve_string_to_delete_if_present)
            if parameters_string in self.file:
                del self.file[parameters_string]
            if curves_string in self.file:
                del self.file[curves_string]
            self.close()
            print(("Stopping curve generation due to KeyboardInterrupt at " +\
                "{!s}.").format(time.ctime()))
    
    def get_training_set(self, return_parameters=False, return_model=False,\
        return_prior_set=False):
        """
        Gets the (assumed already generated) training set in the file of this
        TrainingSetCreator.
        
        return_model: if True, model is also returned
        return_prior_set: if True, prior_set is also returned
        
        returns: numpy.ndarray whose shape is (num_curves, num_channels),
                 followed by parameters, model and/or prior_set if applicable
        """
        group = self.file['curves']
        de_facto_num_curves = 0
        while 'curve_{:d}'.format(de_facto_num_curves) in group:
            de_facto_num_curves += 1
        if de_facto_num_curves == 0:
            raise RuntimeError("No curves have been generated yet.")
        else:
            num_channels = np.max([group['curve_{:d}'.format(icurve)].size\
                for icurve in range(de_facto_num_curves)])
        training_set = np.ndarray((de_facto_num_curves, num_channels))
        to_keep = []
        for icurve in range(de_facto_num_curves):
            curve = get_hdf5_value(group['curve_{:d}'.format(icurve)])
            if np.any(np.isnan(curve)):
                training_set[icurve,:] = np.nan
            else:
                training_set[icurve,:] = curve
                to_keep.append(icurve)
        to_keep = np.array(to_keep)
        self.close()
        return_value = [training_set[to_keep,:]]
        if return_parameters:
            parameters =\
                np.ndarray((de_facto_num_curves, self.model.num_parameters))
            group = self.file['parameters']
            for icurve in range(de_facto_num_curves):
                parameters[icurve,:] =\
                    get_hdf5_value(group['curve_{:d}'.format(icurve)])
            parameters = parameters[to_keep,:]
            parameters = {param: parameters[:,iparam]\
                for (iparam, param) in enumerate(self.model.parameters)}
            return_value.append(parameters)
            self.close()
        if return_model:
            return_value.append(self.model)
        if return_prior_set:
            return_value.append(self.prior_set)
        if len(return_value) == 1:
            return return_value[0]
        else:
            return tuple(return_value)
    
    def get_bad_parameters(self):
        """
        Gets the bad parameters which were encountered by this
        TrainingSetCreator (i.e. the parameter values which produce curves that
        throw errors or return nan's)
        
        returns: (model, parameters) where model is a Model object and
                 parameters is a dictionary of arrays indexed by parameter name
        """
        group = self.file['curves']
        de_facto_num_curves = 0
        while 'curve_{:d}'.format(de_facto_num_curves) in group:
            de_facto_num_curves += 1
        to_keep = []
        for icurve in range(de_facto_num_curves):
            temp_curve_string = 'curve_{:d}'.format(icurve)
            if np.any(np.isnan(get_hdf5_value(group[temp_curve_string]))):
                to_keep.append(icurve)
        self.close()
        parameters = np.ndarray((len(to_keep), self.model.num_parameters))
        group = self.file['parameters']
        for (iicurve, icurve) in enumerate(to_keep):
            parameters[iicurve,:] =\
                get_hdf5_value(group['curve_{:d}'.format(icurve)])
        parameters = {param: parameters[:,iparam]\
            for (iparam, param) in enumerate(self.model.parameters)}
        return (self.model, parameters)
    
    @staticmethod
    def load_training_set(file_name, return_parameters=False,\
        return_model=False, return_prior_set=False):
        """
        Loads a training set from the TrainingSetCreator which was saved to
        the given file name.
        
        file_name: the file in which a TrainingSetCreator was saved
        return_parameters: if True, dictionary of parameter values is also
                           returned
        return_model: if True, model is also returned
        return_prior_set: if True, prior_set is also returned
        
        returns: numpy.ndarray whose shape is (num_curves, num_channels),
                 followed by model and/or prior_set if applicable
        """
        hdf5_file = h5py.File(file_name, 'r')
        group = hdf5_file['curves']
        num_curves = 0
        while 'curve_{:d}'.format(num_curves) in group:
            num_curves += 1
        if num_curves == 0:
            raise RuntimeError("No curves have been stored in the given file.")
        else:
            num_channels = np.max([group['curve_{:d}'.format(icurve)].size\
                for icurve in range(num_curves)])
        training_set = np.ndarray((num_curves, num_channels))
        to_keep = []
        for icurve in range(num_curves):
            curve = get_hdf5_value(group['curve_{:d}'.format(icurve)])
            if np.any(np.isnan(curve)):
                training_set[icurve,:] = np.nan
            else:
                training_set[icurve,:] = curve
                to_keep.append(icurve)
        to_keep = np.array(to_keep)
        return_value = [training_set[to_keep,:]]
        model = load_model_from_hdf5_group(hdf5_file['model'])
        if return_parameters:
            parameters = np.ndarray((num_curves, model.num_parameters))
            group = hdf5_file['parameters']
            for icurve in range(num_curves):
                parameters[icurve,:] =\
                    get_hdf5_value(group['curve_{:d}'.format(icurve)])
            parameters = parameters[to_keep,:]
            parameters = {parameter: parameters[:,iparameter]\
                for (iparameter, parameter) in enumerate(model.parameters)}
            return_value.append(parameters)
        if return_model:
            return_value.append(model)
        if return_prior_set:
            prior_set =\
                DistributionSet.load_from_hdf5_group(hdf5_file['prior_set'])
            return_value.append(prior_set)
        hdf5_file.close()
        if len(return_value) == 1:
            return return_value[0]
        else:
            return tuple(return_value)
    
    @staticmethod
    def load_training_sets(file_names, return_parameters=False,\
        return_model=False):
        """
        Loads training sets from the multiple TrainingSetCreator objects which
        were saved to the given file names.
        
        file_names: list of files in which TrainingSetCreator object(s) were
                    saved
        return_parameters: if True, dictionary of parameter values is also
                           returned
        return_model: if True, model is also returned
        
        returns: numpy.ndarray whose shape is (total_num_curves, num_channels),
                 followed by parameters and model if applicable
        """
        if len(file_names) == 0:
            raise ValueError("file_names must be a non-empty list of file " +\
                "names at which TrainingSetCreator(s) were saved.")
        training_set_chunks = []
        if return_parameters:
            parameters_chunks = []
        models = []
        for file_name in file_names:
            load = TrainingSetCreator.load_training_set(file_name,\
                return_parameters=return_parameters, return_model=True,\
                return_prior_set=False)
            training_set_chunks.append(load[0])
            if return_parameters:
                parameters_chunks.append(load[1])
            models.append(load[-1])
        if not all([(model == models[0]) for model in models]):
            raise ValueError("The models used for the different " +\
                "TrainingSetCreator objects were not identical.")
        model = models[0]
        training_set = np.concatenate(training_set_chunks, axis=0)
        return_value = [training_set]
        if return_parameters:
            parameters =\
                {parameter: np.concatenate([parameters_chunk[parameter]\
                for parameters_chunk in parameters_chunks])\
                for parameter in model.parameters}
            return_value.append(parameters)
        if return_model:
            return_value.append(model)
        if len(return_value) == 1:
            return return_value[0]
        else:
            return tuple(return_value)
    
    @staticmethod
    def load_bad_parameters(file_name):
        """
        Gets the bad parameters which were encountered by this
        TrainingSetCreator (i.e. the parameter values which produce curves that
        throw errors or return nan's)
        
        returns: (model, parameters) where model is a Model object and
                 parameters is a dictionary of arrays indexed by parameter name
        """
        hdf5_file = h5py.File(file_name, 'r')
        group = hdf5_file['curves']
        num_curves = 0
        while 'curve_{:d}'.format(num_curves) in group:
            num_curves += 1
        to_keep = []
        for icurve in range(num_curves):
            temp_curve_string = 'curve_{:d}'.format(icurve)
            if np.any(np.isnan(get_hdf5_value(group[temp_curve_string]))):
                to_keep.append(icurve)
        model = load_model_from_hdf5_group(hdf5_file['model'])
        parameters = np.ndarray((len(to_keep), model.num_parameters))
        group = hdf5_file['parameters']
        for (iicurve, icurve) in enumerate(to_keep):
            parameters[iicurve,:] =\
                get_hdf5_value(group['curve_{:d}'.format(icurve)])
        parameters = {parameter: parameters[:,iparameter]\
            for (iparameter, parameter) in enumerate(model.parameters)}
        hdf5_file.close()
        return (model, parameters)
    
    def close(self):
        """
        Closes the file containing the curves made by this object.
        """
        if hasattr(self, '_file'):
            self.file.close()
            del self._file

