from __future__ import absolute_import
import torch
import numpy as np
import copy
import multiprocessing
from pysurvival import HAS_GPU
from pysurvival import utils
from pysurvival.utils import neural_networks as nn
from pysurvival.utils import optimization as opt
from pysurvival.models import BaseModel
import sys


# %matplotlib inline

class BaseMultiTaskModel(BaseModel):
    """ Base class for all  Multi-Task estimators:
        * Multi-Task Logistic Regression model (MTLR)
        * Neural Multi-Task Logistic Regression model (N-MTLR)
    BaseMultiTaskModel shouldn't be used as is. 
    The underlying model is written in PyTorch.

    The original Multi-Task model, a.k.a the Multi-Task Logistic Regression
    model (MTLR), was first introduced by  Chun-Nam Yu et al. in 
    *Learning Patient-Specific Cancer Survival Distributions as a Sequence of 
    Dependent Regressors*

    The Neural Multi-Task Logistic Regression model (N-MTLR) was developed 
    by S. Fotso in the paper *Deep Neural Networks for Survival 
    Analysis Based on a Multi-Task Framework*, allowing the use of 
    Neural Networks within the original design.

    Parameters
    ----------
    * `structure`:  **list of dictionaries** -- 
        Provides the structure of the MLP built within the N-MTLR.
        
        ex: `structure = [ {'activation': 'ReLU', 'num_units': 128}, ]`.

        Each dictionary corresponds to a fully connected hidden layer:

        * `num_units` is the number of hidden units in this layer
        * `activation` is the activation function that will be used. 
        The list of all available activation functions can be found :
            * Atan
            * BentIdentity
            * BipolarSigmoid
            * CosReLU
            * ELU
            * Gaussian
            * Hardtanh
            * Identity
            * InverseSqrt
            * LeakyReLU
            * LeCunTanh
            * LogLog
            * LogSigmoid
            * ReLU
            * SELU
            * Sigmoid
            * Sinc
            * SinReLU
            * Softmax
            * Softplus
            * Softsign
            * Swish
            * Tanh
    
        In case there are more than one dictionary, 
        each hidden layer will be applied in the resulting MLP, 
        using the order it is provided in the structure:
        ex: structure = [ {'activation': 'relu', 'num_units': 128}, 
                          {'activation': 'tanh', 'num_units': 128}, ] 

    * `bins`: **int** *(default=100)* -- 
         Number of subdivisions of the time axis 

    * `auto_scaler`: **boolean** *(default=True)* -- 
        Determines whether a sklearn scaler should be automatically applied
    """

    def __init__(self, structure, bins=100, auto_scaler=True):

        # Saving the attributes
        self.loss_values = []
        self.bins = bins
        self.structure = structure

        # Initializing the elements from BaseModel
        super(BaseMultiTaskModel, self).__init__(auto_scaler)

    def get_times(self, T, is_min_time_zero=True, extra_pct_time=0.1):
        """ Building the time axis (self.times) as well as the time intervals 
            ( all the [ t(k-1), t(k) ) in the time axis.
        """

        # Setting the min_time and max_time
        max_time = max(T)
        if is_min_time_zero:
            min_time = 0.
        else:
            min_time = min(T)

        # Setting optional extra percentage time
        if 0. <= extra_pct_time <= 1.:
            p = extra_pct_time
        else:
            raise Exception("extra_pct_time has to be between [0, 1].")

            # Building time points and time buckets
        self.times = np.linspace(min_time, max_time * (1. + p), self.bins)
        self.get_time_buckets()
        self.num_times = len(self.time_buckets)

    def compute_XY(self, X, T, E, is_min_time_zero, extra_pct_time):
        """ Given the survival_times, events and time_points vectors, 
            it returns a ndarray of the encodings for all units 
            such that:
                Y = [[0, 0, 1, 0, 0],  # unit experienced an event at t = 3
                     [0, 1, 0, 0, 0],  # unit experienced an event at t = 2
                     [0, 1, 1, 1, 1],] # unit was censored at t = 2
        """

        # building times axis
        self.get_times(T, is_min_time_zero, extra_pct_time)
        n_units = T.shape[0]

        # Initializing the output variable
        Y_cens, Y_uncens = [], []
        X_cens, X_uncens = [], []

        if isinstance(X, list):
            for input_ in X:
                X_cens.append([])
                X_uncens.append([])

        # Building the output variable
        for i, (t, e) in enumerate(zip(T, E)):
            y = np.zeros(self.num_times + 1)
            min_abs_value = [abs(a_j_1 - t) for (a_j_1, a_j) in self.time_buckets]
            index = np.argmin(min_abs_value)

            if e == 1:
                y[index] = 1.
                if isinstance(X, list):
                    for j, input_ in enumerate(X):
                        X_uncens[j].append(input_[i, :].tolist())
                else:
                    X_uncens.append(X[i, :].tolist())
                Y_uncens.append(y.tolist())

            else:
                y[(index):] = 1.

                if isinstance(X, list):
                    for j, input_ in enumerate(X):
                        X_cens[j].append(input_[i, :].tolist())
                else:
                    X_cens.append(X[i, :].tolist())
                Y_cens.append(y.tolist())

        # Transform into torch.Tensor
        if isinstance(X, list):
            for j, input_ in enumerate(X_cens):
                X_cens[j] = torch.FloatTensor(input_)
                if torch.cuda.is_available():
                    X_cens[j] = X_cens[j].cuda()
            for j, input_ in enumerate(X_uncens):
                X_uncens[j] = torch.FloatTensor(input_)
                if torch.cuda.is_available():
                    X_uncens[j] = X_uncens[j].cuda()
        else:
            X_cens = torch.FloatTensor(X_cens)
            X_uncens = torch.FloatTensor(X_uncens)
            if torch.cuda.is_available():
                X_cens = X_cens.cuda()
                X_uncens = X_uncens.cuda()

        Y_cens = torch.FloatTensor(Y_cens)
        Y_uncens = torch.FloatTensor(Y_uncens)
        if torch.cuda.is_available():
            Y_cens = Y_cens.cuda()
            Y_uncens = Y_uncens.cuda()

        return X_cens, X_uncens, Y_cens, Y_uncens

    def loss_function(self, model, X_cens, X_uncens, Y_cens, Y_uncens,
                      Triangle, l2_reg, l2_smooth, min_clamp_value=1e-8,
                      max_clamp_value=torch.finfo(torch.float32).max-1):
        """ Computes the loss function of the any MTLR model. 
            All the operations have been vectorized to ensure optimal speed
        """

        score_cens = model(X_cens)
        score_uncens = model(X_uncens)
        # Likelihood Calculations -- Uncensored
        temp = torch.clamp(torch.mm(score_uncens, Triangle), min=1e-8, max=88.5)
        phi_uncens = torch.clamp(torch.exp(temp), max=max_clamp_value)
        reduc_phi_uncens = torch.sum(phi_uncens * Y_uncens, dim=1)

        # Likelihood Calculations -- Censored
        temp = torch.clamp(torch.mm(score_cens, Triangle), min=1e-8, max=88.5)
        phi_cens = torch.clamp(torch.exp(temp), max=max_clamp_value)
        reduc_phi_cens = torch.sum(phi_cens * Y_cens, dim=1)

        # Likelihood Calculations -- Normalization
        temp = torch.clamp(torch.mm(score_uncens, Triangle), min=1e-8, max=88.5)
        z_uncens = torch.clamp(torch.exp(temp), max=max_clamp_value)
        reduc_z_uncens = torch.sum(z_uncens, dim=1)

        temp = torch.clamp(torch.mm(score_cens, Triangle), min=1e-8, max=88.5)
        z_cens = torch.clamp(torch.exp(temp), max=max_clamp_value)
        reduc_z_cens = torch.sum(z_cens, dim=1)

        reduc_phi_uncens = torch.clamp(reduc_phi_uncens, min=min_clamp_value, max=max_clamp_value)
        reduc_phi_cens = torch.clamp(reduc_phi_cens, min=min_clamp_value, max=max_clamp_value)
        reduc_z_uncens = torch.clamp(reduc_z_uncens, min=min_clamp_value, max=max_clamp_value)
        reduc_z_cens = torch.clamp(reduc_z_cens, min=min_clamp_value, max=max_clamp_value)

        # MTLR cost function
        loss = - (
                torch.sum(torch.log(reduc_phi_uncens))
                + torch.sum(torch.log(reduc_phi_cens))

                - torch.sum(torch.log(reduc_z_uncens))
                - torch.sum(torch.log(reduc_z_cens))
        )

        # print(f"loss_loss: {loss}")
        # Adding the regularized loss
        nb_set_parameters = len(list(model.parameters()))
        for i, w in enumerate(model.parameters()):
            loss += l2_reg * torch.sum(w * w) / 2.

            if i >= nb_set_parameters - 2:
                loss += l2_smooth * norm_diff(w)

        return loss

    def fit(self, X, T, E, init_method='glorot_uniform', optimizer='adam',
            lr=1e-4, num_epochs=1000, dropout=0.2, l2_reg=1e-2,
            l2_smooth=1e-2, batch_normalization=False, bn_and_dropout=False,
            verbose=True, extra_pct_time=0.1, is_min_time_zero=True, max_norm=1.0,
            min_clamp_value=1e-8, max_clamp_value=torch.finfo(torch.float32).max -1):
        """ Fit the estimator based on the given parameters.

        Parameters:
        -----------
        * `X` : **array-like**, *shape=(n_samples, n_features)* --
            The input samples.

        * `T` : **array-like** -- 
            The target values describing when the event of interest or censoring
            occurred.

        * `E` : **array-like** --
            The values that indicate if the event of interest occurred i.e.: 
            E[i]=1 corresponds to an event, and E[i] = 0 means censoring, 
            for all i.

        * `init_method` : **str** *(default = 'glorot_uniform')* -- 
            Initialization method to use. Here are the possible options:

            * `glorot_uniform`: Glorot/Xavier uniform initializer
            * `he_uniform`: He uniform variance scaling initializer
            * `uniform`: Initializing tensors with uniform (-1, 1) distribution
            * `glorot_normal`: Glorot normal initializer,
            * `he_normal`: He normal initializer.
            * `normal`: Initializing tensors with standard normal distribution
            * `ones`: Initializing tensors to 1
            * `zeros`: Initializing tensors to 0
            * `orthogonal`: Initializing tensors with a orthogonal matrix,

        * `optimizer`:  **str** *(default = 'adam')* -- 
            iterative method for optimizing a differentiable objective function.
            Here are the possible options:

            - `adadelta`
            - `adagrad`
            - `adam`
            - `adamax`
            - `rmsprop`
            - `sparseadam`
            - `sgd`

        * `lr`: **float** *(default=1e-4)* -- 
            learning rate used in the optimization

        * `num_epochs`: **int** *(default=1000)* -- 
            The number of iterations in the optimization

        * `dropout`: **float** *(default=0.5)* -- 
            Randomly sets a fraction rate of input units to 0 
            at each update during training time, which helps prevent overfitting.

        * `l2_reg`: **float** *(default=1e-4)* -- 
            L2 regularization parameter for the model coefficients

        * `l2_smooth`: **float** *(default=1e-4)* -- 
            Second L2 regularizer that ensures the parameters vary smoothly 
            across consecutive time points.

        * `batch_normalization`: **bool** *(default=True)* -- 
            Applying Batch Normalization or not

        * `bn_and_dropout`: **bool** *(default=False)* -- 
            Applying Batch Normalization and Dropout at the same time

        * `display_loss`: **bool** *(default=True)* -- 
            Whether or not showing the loss function values at each update

        * `verbose`: **bool** *(default=True)* -- 
            Whether or not producing detailed logging about the modeling

        * `extra_pct_time`: **float** *(default=0.1)* -- 
            Providing an extra fraction of time in the time axis

        * `is_min_time_zero`: **bool** *(default=True)* -- 
            Whether the the time axis starts at 0

        * `max_norm`: **float** *(default=1.0)* --
            Max l2 norm for gradient clipping

        **Returns:**

        * self : object


        Example:
        --------
            
        #### 1 - Importing packages
        import numpy as np
        import pandas as pd
        from matplotlib import pyplot as plt
        from sklearn.model_selection import train_test_split
        from pysurvival.models.simulations import SimulationModel
        from pysurvival.models.multi_task import LinearMultiTaskModel
        from pysurvival.utils.metrics import concordance_index
        #%matplotlib inline  # To use with Jupyter notebooks


        #### 2 - Generating the dataset from a Weibull parametric model
        # Initializing the simulation model
        sim = SimulationModel( survival_distribution = 'Weibull',  
                               risk_type = 'linear',
                               censored_parameter = 10.0, 
                               alpha = .01, beta = 3.0 )

        # Generating N random samples 
        N = 1000
        dataset = sim.generate_data(num_samples = N, num_features = 3)

        # Showing a few data-points 
        time_column = 'time'
        event_column = 'event'
        dataset.head(2)

        #### 3 - Creating the modeling dataset
        # Defining the features
        features = sim.features

        # Building training and testing sets #
        index_train, index_test = train_test_split( range(N), test_size = 0.2)
        data_train = dataset.loc[index_train].reset_index( drop = True )
        data_test  = dataset.loc[index_test].reset_index( drop = True )

        # Creating the X, T and E input
        X_train, X_test = data_train[features], data_test[features]
        T_train, T_test = data_train['time'].values, data_test['time'].values
        E_train, E_test = data_train['event'].values, data_test['event'].values

        #### 4 - Initializing a MTLR model and fitting the data.
        # Building a Linear model
        mtlr = LinearMultiTaskModel(bins=50) 
        mtlr.fit(X_train, T_train, E_train, lr=5e-3, init_method='orthogonal')

        # Building a Neural MTLR
        # structure = [ {'activation': 'Swish', 'num_units': 150},  ]
        # mtlr = NeuralMultiTaskModel(structure=structure, bins=150) 
        # mtlr.fit(X_train, T_train, E_train, lr=5e-3, init_method='adam')

        #### 5 - Cross Validation / Model Performances
        c_index = concordance_index(mtlr, X_test, T_test, E_test) #0.95
        print('C-index: {:.2f}'.format(c_index))

        """

        # Checking data format (i.e.: transforming into numpy array)
        X, T, E = utils.check_data(X, T, E)

        input_shape = []
        # Extracting data parameters
        if isinstance(X, list):
            nb_inputs = len(X)
            for data in X:
                nb_units, num_vars = data.shape
                input_shape.append(num_vars)
            # Scaling data
            if self.auto_scaler:
                for index, data in enumerate(X):
                    X[index] = self.scaler.fit_transform(data)
        else:
            nb_inputs = 1
            nb_units, self.num_vars = X.shape
            input_shape.append(self.num_vars)
            # Scaling data
            if self.auto_scaler:
                X = self.scaler.fit_transform(X)

        # Building the time axis, time buckets and output Y
        X_cens, X_uncens, Y_cens, Y_uncens \
            = self.compute_XY(X, T, E, is_min_time_zero, extra_pct_time)

        # Initializing the model
        model = nn.NeuralNet(input_shape, self.num_times, self.structure,
                             init_method, dropout, batch_normalization,
                             bn_and_dropout)

        # Creating the Triangular matrix
        Triangle = np.tri(self.num_times, self.num_times + 1, dtype=np.float32)
        Triangle = torch.FloatTensor(Triangle)

        if torch.cuda.is_available():
            model = model.cuda()
            Triangle = Triangle.cuda()

        # Performing order 1 optimization
        model, loss_values = opt.optimize(self.loss_function, model, optimizer,
                                          lr, num_epochs, verbose, X_cens=X_cens, X_uncens=X_uncens,
                                          Y_cens=Y_cens, Y_uncens=Y_uncens, Triangle=Triangle,
                                          l2_reg=l2_reg, l2_smooth=l2_smooth, max_norm=max_norm,
                                          min_clamp_value=min_clamp_value, max_clamp_value=max_clamp_value)

        # Saving attributes
        self.model = model.eval()
        self.loss_values = loss_values

        return self

    def predict(self, x, t=None):
        """ Predicting the hazard, density and survival functions
        
        Parameters:
        ----------
        * `x` : **array-like** *shape=(n_samples, n_features)* --
            array-like representing the datapoints. 
            x should not be standardized before, the model
            will take care of it

        * `t`: **double** *(default=None)* --
             time at which the prediction should be performed. 
             If None, then return the function for all available t.
        """

        # Convert x into the right format
        x = utils.check_data(x)

        # Scaling the data
        if self.auto_scaler:
            if x.ndim == 1:
                x = self.scaler.transform(x.reshape(1, -1))
            elif x.ndim == 2:
                x = self.scaler.transform(x)
            x = torch.FloatTensor(x)
            if torch.cuda.is_available():
                x = x.cuda()
        else:
            # Ensuring x has 2 dimensions
            if isinstance(x, list):
                for index in range(len(x)):
                    if x[index].ndim == 1:
                        x[index] = np.reshape(x[index], (1, -1))
                    # Transforming into pytorch objects
                    x[index] = torch.FloatTensor(x[index])
                    if torch.cuda.is_available():
                        x[index] = x[index].cuda()
            else:
                if x.ndim == 1:
                    x = np.reshape(x, (1, -1))
                # Transforming into pytorch objects
                x = torch.FloatTensor(x)
                if torch.cuda.is_available():
                    x = x.cuda()

        # Predicting using linear/nonlinear function
        score_torch = self.model(x)
        score = score_torch.data.cpu().numpy()

        # Cretaing the time triangles
        Triangle1 = np.tri(self.num_times, self.num_times + 1)
        Triangle2 = np.tri(self.num_times + 1, self.num_times + 1)

        # Calculating the score, density, hazard and Survival
        phi = np.exp(np.dot(score, Triangle1))
        div = np.repeat(np.sum(phi, 1).reshape(-1, 1), phi.shape[1], axis=1)
        density = (phi / div)
        Survival = np.dot(density, Triangle2)
        hazard = density[:, :-1] / Survival[:, 1:]

        # Returning the full functions of just one time point
        if t is None:
            return hazard, density, Survival
        else:
            min_abs_value = [abs(a_j_1 - t) for (a_j_1, a_j) in self.time_buckets]
            index = np.argmin(min_abs_value)
            return hazard[:, index], density[:, index], Survival[:, index]

    def predict_risk(self, x, use_log=False):
        """ Computing the risk score 

        Parameters:
        -----------
        * `x` : **array-like** *shape=(n_samples, n_features)* --
            array-like representing the datapoints. 
            x should not be standardized before, the model
            will take care of it

        * `use_log`: **bool** *(default=True)* -- 
            Applies the log function to the risk values

        """

        risk = super(BaseMultiTaskModel, self).predict_risk(x)
        if use_log:
            return np.log(risk)
        else:
            return risk


class LinearMultiTaskModel(BaseMultiTaskModel):
    """ LinearMultiTaskModel is the original Multi-Task model, 
        a.k.a the Multi-Task Logistic Regression model (MTLR).
        It was first introduced by  Chun-Nam Yu et al. in 
        Learning Patient-Specific Cancer Survival Distributions 
        as a Sequence of Dependent Regressors
        
        Reference:
        ----------
            * http://www.cs.cornell.edu/~cnyu/papers/nips11_survival.pdf
        
        Parameters:
        ----------
            * bins: int 
                Number of subdivisions of the time axis 

            * auto_scaler: boolean (default=True)
                Determines whether a sklearn scaler should be automatically 
                applied

    """

    def __init__(self, bins=100, auto_scaler=True):
        super(LinearMultiTaskModel, self).__init__(
            structure=None, bins=bins, auto_scaler=auto_scaler)

    def fit(self, X, T, E, init_method='glorot_uniform', optimizer='adam',
            lr=1e-4, num_epochs=1000, l2_reg=1e-2, l2_smooth=1e-2,
            verbose=True, extra_pct_time=0.1, is_min_time_zero=True, max_norm=1.0,
            min_clamp_value=1e-8, max_clamp_value=torch.finfo(torch.float32).max-1):
        super(LinearMultiTaskModel, self).fit(X=X, T=T, E=E,
                                              init_method=init_method, optimizer=optimizer,
                                              lr=lr, num_epochs=num_epochs, dropout=None, l2_reg=l2_reg,
                                              l2_smooth=l2_smooth, batch_normalization=False,
                                              bn_and_dropout=False, verbose=verbose,
                                              extra_pct_time=extra_pct_time, is_min_time_zero=is_min_time_zero,
                                              max_norm=max_norm,
                                              min_clamp_value=min_clamp_value, max_clamp_value=max_clamp_value)

        return self


class NeuralMultiTaskModel(BaseMultiTaskModel):
    """ NeuralMultiTaskModel is the Neural Multi-Task Logistic Regression 
        model (N-MTLR) was developed by Fotso S. in 
        Deep Neural Networks for Survival Analysis Based on a 
        Multi-Task Framework, 
        allowing the use of Neural Networks within the original design.
        
    Reference:
    ----------
    * https://arxiv.org/pdf/1801.05512

    Parameters:
    ----------
    * `structure`:  **list of dictionaries** -- 
        Provides the structure of the MLP built within the N-MTLR.
        
        ex: `structure = [ {'activation': 'ReLU', 'num_units': 128}, ]`.

        Each dictionary corresponds to a fully connected hidden layer:

        * `units` is the number of hidden units in this layer
        * `activation` is the activation function that will be used. 
        The list of all available activation functions can be found :
            * Atan
            * BentIdentity
            * BipolarSigmoid
            * CosReLU
            * ELU
            * Gaussian
            * Hardtanh
            * Identity
            * InverseSqrt
            * LeakyReLU
            * LeCunTanh
            * LogLog
            * LogSigmoid
            * ReLU
            * SELU
            * Sigmoid
            * Sinc
            * SinReLU
            * Softmax
            * Softplus
            * Softsign
            * Swish
            * Tanh
    
        In case there are more than one dictionary, 
        each hidden layer will be applied in the resulting MLP, 
        using the order it is provided in the structure:
        ex: structure = [ {'activation': 'relu', 'num_units': 128}, 
                          {'activation': 'tanh', 'num_units': 128}, ] 

    * `bins`: **int** *(default=100)* -- 
         Number of subdivisions of the time axis 

    * `auto_scaler`: **boolean** *(default=True)* -- 
        Determines whether a sklearn scaler should be automatically applied
    """

    def __init__(self, structure, bins=100, auto_scaler=True):

        # Checking the validity of structure
        structure = nn.check_mlp_structure(structure)

        # print(structure)

        # Initializing the instance
        super(NeuralMultiTaskModel, self).__init__(
            structure=structure, bins=bins, auto_scaler=auto_scaler)

    def __repr__(self):
        """ Representing the class object """

        if self.structure is None:
            super(NeuralMultiTaskModel, self).__repr__()
            return self.name

        else:
            S = len(self.structure)
            self.name = self.__class__.__name__
            empty = len(self.name)
            self.name += '( '
            for i, s in enumerate(self.structure):
                if isinstance(s, list):
                    for s_ in s:
                        n = 'Layer({}): '.format(i + 1)
                        activation = nn.activation_function(s_['activation'],
                                                            return_text=True)
                        n += 'activation = {}, '.format(s_['activation'])
                        if 'num_units' in s_.keys():
                            n += 'units = {} '.format(s_['num_units'])

                        if i != S - 1:
                            self.name += n + '; \n'
                            self.name += empty * ' ' + '  '
                        else:
                            self.name += n
                    self.name = self.name + ')'
                else:
                    n = 'Layer({}): '.format(i + 1)
                    activation = nn.activation_function(s['activation'],
                                                        return_text=True)
                    n += 'activation = {}, '.format(s['activation'])
                    if 'num_units' in s.keys():
                        n += 'units = {} '.format(s['num_units'])

                    if i != S - 1:
                        self.name += n + '; \n'
                        self.name += empty * ' ' + '  '
                    else:
                        self.name += n
                self.name = self.name + ')'

            return self.name


def norm_diff(W):
    """ Special norm function for the last layer of the MTLR """
    dims = len(W.shape)
    if dims == 1:
        diff = W[1:] - W[:-1]
    elif dims == 2:
        diff = W[1:, :] - W[:-1, :]
    return torch.sum(diff * diff)
