# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals

import warnings

import deepiv.samplers as samplers
import deepiv.densities as densities

from deepiv.custom_gradients import replace_gradients_mse

from keras.models import Model
from keras import backend as K
from keras.layers import Lambda, InputLayer
from keras.engine import topology
try:
    import h5py
except ImportError:
    h5py = None


import keras.utils

import numpy
from sklearn import linear_model
from sklearn.decomposition import PCA
from scipy.stats import norm

class Treatment(Model):
    '''
    Adds sampling functionality to a Keras model and extends the losses to support
    mixture of gaussian losses.

    # Argument
    '''

    def _get_sampler_by_string(self, loss):
        output = self.outputs[0]
        inputs = self.inputs

        if loss in ["MSE", "mse", "mean_squared_error"]:
            output += samplers.random_normal(K.shape(output), mean=0.0, std=1.0)
            draw_sample = K.function(inputs + [K.learning_phase()], [output])

            def sample_gaussian(inputs, use_dropout=False):
                '''
                Helper to draw samples from a gaussian distribution
                '''
                return draw_sample(inputs + [int(use_dropout)])[0]

            return sample_gaussian

        elif loss == "binary_crossentropy":
            output = K.random_binomial(K.shape(output), p=output)
            draw_sample = K.function(inputs + [K.learning_phase()], [output])

            def sample_binomial(inputs, use_dropout=False):
                '''
                Helper to draw samples from a binomial distribution
                '''
                return draw_sample(inputs + [int(use_dropout)])[0]

            return sample_binomial

        elif loss in ["mean_absolute_error", "mae", "MAE"]:
            output += samplers.random_laplace(K.shape(output), mu=0.0, b=1.0)
            draw_sample = K.function(inputs + [K.learning_phase()], [output])
            def sample_laplace(inputs, use_dropout=False):
                '''
                Helper to draw samples from a Laplacian distribution
                '''
                return draw_sample(inputs + [int(use_dropout)])[0]

            return sample_laplace

        elif loss == "mixture_of_gaussians":
            pi, mu, log_sig = densities.split_mixture_of_gaussians(output, self.n_components)
            samples = samplers.random_gmm(pi, mu, K.exp(log_sig))
            draw_sample = K.function(inputs + [K.learning_phase()], [samples])
            return lambda inputs, use_dropout: draw_sample(inputs + [int(use_dropout)])[0]

        else:
            raise NotImplementedError("Unrecognised loss: %s.\
                                       Cannot build a generic sampler" % loss)

    def _prepare_sampler(self, loss):
        '''
        Build sampler
        '''
        if isinstance(loss, str):
            self._sampler = self._get_sampler_by_string(loss)
        else:
            warnings.warn("You're using a custom loss function. Make sure you implement\
                           the model's sample() fuction yourself.")

    def compile(self, optimizer, loss, metrics=None, loss_weights=None,
                sample_weight_mode=None, n_components=None, **kwargs):
        '''
        Overrides the existing keras compile function to add a sampler building
        step to the model compilation phase. Once compiled, one can draw samples
        from the network using the sample() function and adds support for mixture
        of gaussian loss.

        '''
        if loss == "mixture_of_gaussians":
            if n_components is None:
                raise Exception("When using mixture of gaussian loss you must\
                                 supply n_components argument")
            self.n_components = n_components
            self._prepare_sampler(loss)
            loss = lambda y_true, y_pred: densities.mixture_of_gaussian_loss(y_true,
                                                                             y_pred,
                                                                             n_components)

            def predict_mean(x, batch_size=32, verbose=0):
                '''
                Helper to just predict the expected value of the mixture
                of gaussian rather than the parameters for the distribution.
                '''
                y_hat = super(Treatment, self).predict(x, batch_size, verbose)
                n_c = n_components
                return (y_hat[:, 0:n_c] * y_hat[:, n_c:2*n_c]).sum(axis=1, keepdims=True)

            self.predict_mean = predict_mean
        else:
            self._prepare_sampler(loss)

        super(Treatment, self).compile(optimizer, loss, metrics=metrics, loss_weights=loss_weights,
                                       sample_weight_mode=sample_weight_mode, **kwargs)

    def sample(self, inputs, n_samples=1, use_dropout=False):
        '''
        Draw samples from the keras model.
        '''
        if hasattr(self, "_sampler"):
            if not isinstance(inputs, list):
                inputs = [inputs]
            inputs = [i.repeat(n_samples, axis=0) for i in inputs]
            return self._sampler(inputs, use_dropout)
        else:
            raise Exception("Compile model with loss before sampling")

class Response(Model):
    '''
    Extends the Keras Model class to support sampling from the Treatment
    model during training.

    Overwrites the existing fit_generator function.

    # Arguments
    In addition to the standard model arguments, a Response object takes
    a Treatment object as input so that it can sample from the fitted treatment
    distriubtion during training.
    '''
    def __init__(self, treatment, **kwargs):
        if isinstance(treatment, Treatment):
            self.treatment = treatment
        else:
            raise TypeError("Expected a treatment model of type Treatment. \
                             Got a model of type %s. Remember to train your\
                             treatment model first." % type(treatment))
        super(Response, self).__init__(**kwargs)

    def compile(self, optimizer, loss, metrics=None, loss_weights=None, sample_weight_mode=None,
                unbiased_gradient=False,n_samples=1, batch_size=None):
        super(Response, self).compile(optimizer=optimizer, loss=loss, loss_weights=loss_weights,
                                      sample_weight_mode=sample_weight_mode)
        self.unbiased_gradient = unbiased_gradient
        if unbiased_gradient:
            if loss in ["MSE", "mse", "mean_squared_error"]:
                if batch_size is None:
                    raise ValueError("Must supply a batch_size argument if using unbiased gradients. Currently batch_size is None.")
                replace_gradients_mse(self, optimizer, batch_size=batch_size, n_samples=n_samples)
            else:
                warnings.warn("Unbiased gradient only implemented for mean square error loss. It is unnecessary for\
                              logistic losses and currently not implemented for absolute error losses.")
            

    def fit(self, x=None, y=None, batch_size=512, epochs=1, verbose=1, callbacks=None,
            validation_data=None, class_weight=None, initial_epoch=0, samples_per_batch=None,
            seed=None, observed_treatments=None):
        '''
        Trains the model by sampling from the fitted treament distribution.

        # Arguments
            x: list of numpy arrays. The first element should *always* be the instrument variables.
            y: (numpy array). Target response variables.
            The remainder of the arguments correspond to the Keras definitions.
        '''
        batch_size = numpy.minimum(y.shape[0], batch_size)
        if seed is None:
            seed = numpy.random.randint(0, 1e6)
        if samples_per_batch is None:
            if self.unbiased_gradient:
                samples_per_batch = 2
            else:
                samples_per_batch = 1

        if observed_treatments is None:
            generator = SampledSequence(x[1:], x[0], y, batch_size, self.treatment.sample, samples_per_batch)
        else:
            generator = OnesidedUnbaised(x[1:], x[0], y, observed_treatments, batch_size,
                                         self.treatment.sample, samples_per_batch)
        
        steps_per_epoch = y.shape[0]  // batch_size
        super(Response, self).fit_generator(generator=generator,
                                            steps_per_epoch=steps_per_epoch,
                                            epochs=epochs, verbose=verbose,
                                            callbacks=callbacks, validation_data=validation_data,
                                            class_weight=class_weight, initial_epoch=initial_epoch)

    def fit_generator(self, **kwargs):
        '''
        We use override fit_generator to support sampling from the treatment model during training.

        If you need this functionality, you'll need to build a generator that samples from the
        treatment and performs whatever transformations you're performing. Please submit a pull
        request if you implement this.
        '''
        raise NotImplementedError("We use override fit_generator to support sampling from the\
                                   treatment model during training.")

    def expected_representation(self, x, z, n_samples=100, batch_size=None, seed=None):
        inputs = [z, x]
        if not hasattr(self, "_E_representation"):
            if batch_size is None:
                batch_size = inputs[0].shape[0]
                steps = 1
            else:
                steps = inputs[0].shape[0] // batch_size

            intermediate_layer_model = Model(inputs=self.inputs,
                                             outputs=self.layers[-2].output)
            
            def pred(inputs, n_samples=100, seed=None):
                features = inputs[1]

                samples = self.treatment.sample(inputs, n_samples)
                batch_features = [features.repeat(n_samples, axis=0)] + [samples]
                representation = intermediate_layer_model.predict(batch_features)
                return representation.reshape((inputs[0].shape[0], n_samples, -1)).mean(axis=1)
            self._E_representation = pred
            return self._E_representation(inputs, n_samples, seed)
        else:
            return self._E_representation(inputs, n_samples, seed)

    def conditional_representation(self, x, p):
        inputs = [x, p]
        if not hasattr(self, "_c_representation"):          
            intermediate_layer_model = Model(inputs=self.inputs,
                                             outputs=self.layers[-2].output)

            self._c_representation = intermediate_layer_model.predict
            return self._c_representation(inputs)
        else:
            return self._c_representation(inputs)

    def dropout_predict(self, x, z, n_samples=100):
        if isinstance(x, list):
            inputs = [z] + x
        else:
            inputs = [z, x]
        if not hasattr(self, "_dropout_predict"):
            
            predict_with_dropout = K.function(self.inputs + [K.learning_phase()],
                                              [self.layers[-1].output])

            def pred(inputs, n_samples = 100):
                # draw samples from the treatment network with dropout turned on
                samples = self.treatment.sample(inputs, n_samples, use_dropout=True)
                # prepare inputs for the response network
                rep_inputs = [i.repeat(n_samples, axis=0) for i in inputs[1:]] + [samples]
                # return outputs from the response network with dropout turned on (learning_phase=0)
                return predict_with_dropout(rep_inputs + [1])[0]
            self._dropout_predict = pred
            return self._dropout_predict(inputs, n_samples)
        else:
            return self._dropout_predict(inputs, n_samples)

    def credible_interval(self, x, z, n_samples=100, p=0.95):
        '''
        Return a credible interval of size p using dropout variational inference.
        '''
        if isinstance(x, list):
            n = x[0].shape[0]
        else:
            n = x.shape[0]
        alpha = (1-p) / 2.
        samples = self.dropout_predict(x, z, n_samples).reshape((n, n_samples, -1))
        upper = numpy.percentile(samples.copy(), 100*(p+alpha), axis=1)
        lower = numpy.percentile(samples.copy(), 100*(alpha), axis=1)
        return lower, upper

    def _add_constant(self, X):
        return numpy.concatenate((numpy.ones((X.shape[0], 1)), X), axis=1)
    
    def predict_confidence(self, x, p):
        if hasattr(self, "_predict_confidence"):
            return self._predict_confidence(x, p)
        else:
            raise Exception("Call fit_confidence_interval before running predict_confidence")

    
    def fit_confidence_interval(self, x_lo, z_lo, p_lo, y_lo, n_samples=100, alpha=0.):
        eta_bar = self.expected_representation(x=x_lo, z=z_lo, n_samples=n_samples)
        pca = PCA(1-1e-16, svd_solver="full", whiten=True)
        pca.fit(eta_bar)

        eta_bar = pca.transform(eta_bar)
        eta_lo_prime = pca.transform(self.conditional_representation(x_lo, p_lo))
        eta_lo = self._add_constant(eta_lo_prime)

        ols1 = linear_model.Ridge(alpha=alpha, fit_intercept=True)
        ols1.fit(eta_bar, eta_lo_prime)
        hhat = ols1.predict(eta_bar)
        ols2 = linear_model.Ridge(alpha=alpha, fit_intercept=False)
        ols2.fit(self._add_constant(hhat), y_lo)

        yhat = ols2.predict(eta_lo)
        hhi = numpy.linalg.inv(numpy.dot(eta_lo.T, eta_lo))

        heh = numpy.dot(eta_lo.T, numpy.square(y_lo - yhat) * eta_lo)
        V = numpy.dot(numpy.dot(hhi, heh), hhi)

        def pred(xx, pp):
            H = self._add_constant(pca.transform(self.conditional_representation(xx,pp)))
            sdhb = numpy.sqrt(numpy.diag(numpy.dot(numpy.dot(H, V), H.T)))
            hb = ols2.predict(H).flatten()
            return hb, sdhb
        
        self._predict_confidence = pred




class SampledSequence(keras.utils.Sequence):
    def __init__(self, features, instruments, outputs, batch_size, sampler, n_samples=1, seed=None):
        self.rng = numpy.random.RandomState(seed)
        if not isinstance(features, list):
            features = [features.copy()]
        else:
            features = [f.copy() for f in features]
        self.features = features
        self.instruments = instruments.copy()
        self.outputs = outputs.copy()
        if batch_size < self.instruments.shape[0]:
            self.batch_size = batch_size
        else:
            self.batch_size = self.instruments.shape[0]
        self.sampler = sampler
        self.n_samples = n_samples
        self.current_index = 0
        self.shuffle()

    def __len__(self):
        if isinstance(self.outputs, list):
            return self.outputs[0].shape[0] // self.batch_size
        else:
            return self.outputs.shape[0] // self.batch_size

    def shuffle(self):
        idx = self.rng.permutation(numpy.arange(self.instruments.shape[0]))
        self.instruments = self.instruments[idx,:]
        self.outputs = self.outputs[idx,:]
        self.features = [f[idx,:] for f in self.features]
    
    def __getitem__(self,idx):
        instruments = [self.instruments[idx*self.batch_size:(idx+1)*self.batch_size, :]]
        features = [inp[idx*self.batch_size:(idx+1)*self.batch_size, :] for inp in self.features]
        sampler_input = instruments + features
        samples = self.sampler(sampler_input, self.n_samples)
        batch_features = [f[idx*self.batch_size:(idx+1)*self.batch_size].repeat(self.n_samples, axis=0) for f in self.features] + [samples]
        batch_y = self.outputs[idx*self.batch_size:(idx+1)*self.batch_size].repeat(self.n_samples, axis=0)
        if idx == (len(self) - 1):
            self.shuffle()
        return batch_features, batch_y

class OnesidedUnbaised(SampledSequence):
    def __init__(self, features, instruments, outputs, treatments, batch_size, sampler, n_samples=1, seed=None):
        self.rng = numpy.random.RandomState(seed)
        if not isinstance(features, list):
            features = [features.copy()]
        else:
            features = [f.copy() for f in features]
        self.features = features
        self.instruments = instruments.copy()
        self.outputs = outputs.copy()
        self.treatments = treatments.copy()
        self.batch_size = batch_size
        self.sampler = sampler
        self.n_samples = n_samples
        self.current_index = 0
        self.shuffle()

    def shuffle(self):
        idx = self.rng.permutation(numpy.arange(self.instruments.shape[0]))
        self.instruments = self.instruments[idx,:]
        self.outputs = self.outputs[idx,:]
        self.features = [f[idx,:] for f in self.features]
        self.treatments = self.treatments[idx,:]

    def __getitem__(self, idx):
        instruments = [self.instruments[idx*self.batch_size:(idx+1)*self.batch_size, :]]
        features = [inp[idx*self.batch_size:(idx+1)*self.batch_size, :] for inp in self.features]
        observed_treatments = self.treatments[idx*self.batch_size:(idx+1)*self.batch_size, :]
        sampler_input = instruments + features
        samples = self.sampler(sampler_input, self.n_samples // 2)
        samples = numpy.concatenate([observed_treatments, samples], axis=0)
        batch_features = [f[idx*self.batch_size:(idx+1)*self.batch_size].repeat(self.n_samples, axis=0) for f in self.features] + [samples]
        batch_y = self.outputs[idx*self.batch_size:(idx+1)*self.batch_size].repeat(self.n_samples, axis=0)
        if idx == (len(self) - 1):
            self.shuffle()
        return batch_features, batch_y

def load_weights(filepath, model):
    if h5py is None:
        raise ImportError('`load_weights` requires h5py.')

    with h5py.File(filepath, mode='r') as f:
        # set weights
        topology.load_weights_from_hdf5_group(f['model_weights'], model.layers)

    return model

