# -*- coding: utf-8 -*-
# @Time    : 25/03/2023
# @Author  : Ing. Jorge Lara
# @Email   : jlara@iee.unsj.edu.ar
# @File    : ------------
# @Software: PyCharm


import torch
import numpy as np
from scipy import optimize
import pandas as pd


class generate_missing_data:

    def __init__(
            self,
            seed: int = 123
    ):
        self.seed = seed

    # Function produce_NA for generating missing values ------------------------------------------------------
    def produce_NaN(
            self,
            df_true: pd.DataFrame = None,
            p_miss: float = 0.40,
            mecha: str = "MCAR",
            opt="logistic",
            p_obs=0.05,
            q=0.3
    ):

        """
        Generate missing values for specifics missing-datasets mechanism and proportion of missing values.

        Parameters
        ----------
        X : torch.DoubleTensor or np.ndarray, shape (n, d)
            Data for which missing values will be simulated.
            If a numpy array is provided, it will be converted to a pytorch tensor.
        p_miss : float
            Proportion of missing values to generate for variables which will have missing values.
        mecha : str,
                Indicates the missing-datasets mechanism to be used. "MCAR" by default, "MAR", "MNAR" or "Mask"
        opt: str,
                For mecha = "MNAR", it indicates how the missing-datasets mechanism is generated: using a logistic regression ("logistic"), quantile censorship ("quantile") or logistic regression for generating a self-masked MNAR mechanism ("selfmasked").
        p_obs : float
                If mecha = "MAR", or mecha = "MNAR" with opt = "logistic" or "quanti", proportion of variables with *no* missing values that will be used for the logistic masking model.
        q : float
            If mecha = "MNAR" and opt = "quanti", quantile level at which the cuts should occur.

        Returns
        ----------
        A dictionnary containing:
        'X_init': the initial datasets matrix.
        'X_incomp': the datasets with the generated missing values.
        'mask': a matrix indexing the generated missing values.s
        :param df_true:
        """

        X = df_true.copy()
        X = X.to_numpy()
        to_torch = torch.is_tensor(X)  ## output a pytorch tensor, or a numpy array
        if not to_torch:
            X = X.astype(np.float32)
            X = torch.from_numpy(X)

        if mecha == "MAR":
            mask = generate_missing_data(self.seed)._MAR_mask(X, p_miss, p_obs).double()
        elif mecha == "MNAR" and opt == "logistic":
            mask = generate_missing_data(self.seed)._MNAR_mask_logistic(X, p_miss, p_obs).double()
        elif mecha == "MNAR" and opt == "quantile":
            mask = generate_missing_data(self.seed)._MNAR_mask_quantiles(X, p_miss, q, 1 - p_obs).double()
        elif mecha == "MNAR" and opt == "selfmasked":
            mask = generate_missing_data(self.seed)._MNAR_self_mask_logistic(X, p_miss).double()
        else:
            mask = (torch.rand(X.shape) < p_miss).double()

        X_nas = X.clone()
        X_nas[mask.bool()] = np.nan
        X_nas.double()


        df_data_miss = pd.DataFrame(X_nas, columns=df_true.columns)
        df_data_miss['datetime'] = df_true.index
        df_data_miss.set_index('datetime', inplace=True)

        return df_data_miss

    def produce_NaN_pattern(self, df_data: pd.DataFrame = None, columns: list = None, prob_missing: float = 0.3, Random: bool = False):

        if columns is None:
            columns = list(df_data.columns)
        df_data = df_data.copy()
        np.random.seed(self.seed)
        W_train = np.random.choice(2, df_data[columns].shape, p=[prob_missing, 1 - prob_missing])
        W_train = pd.DataFrame(W_train, columns=columns).replace(0, np.nan)
        W_train['datetime'] = df_data.index
        W_train.set_index('datetime', inplace=True)
        if not Random:
            aux = list(W_train.columns)
            aux2 = aux[1:]
            for i in aux2:
                W_train[i] = W_train[aux[0]]
        X_train = (df_data[columns] * W_train)
        df_data[columns] = X_train
        return df_data

    def _nanmean(self, v, *args, **kwargs):
        """
        A Pytorch version on Numpy's nanmean
        """
        v = v.clone()
        is_nan = torch.isnan(v)
        v[is_nan] = 0
        return v.sum(*args, **kwargs) / (~is_nan).float().sum(*args, **kwargs)

    #### Quantile ######
    def _quantile(self, X, q, dim=None):
        """
        Returns the q-th quantile.

        Parameters
        ----------
        X : torch.DoubleTensor or torch.cuda.DoubleTensor, shape (n, d)
            Input datasets.

        q : float
            Quantile level (starting from lower values).

        dim : int or None, default = None
            Dimension allong which to compute quantiles. If None, the tensor is flattened and one value is returned.


        Returns
        -------
            quantiles : torch.DoubleTensor

        """
        return X.kthvalue(int(q * len(X)), dim=dim)[0]

    #### Automatic selection of the regularization parameter ####
    def _pick_epsilon(self, X, quant=0.5, mult=0.05, max_points=2000):
        """
            Returns a quantile (times a multiplier) of the halved pairwise squared distances in X.
            Used to select a regularization parameter for Sinkhorn distances.

        Parameters
        ----------
        X : torch.DoubleTensor or torch.cuda.DoubleTensor, shape (n, d)
            Input datasets on which distances will be computed.

        quant : float, default = 0.5
            Quantile to return (default is median).

        mult : float, default = 0.05
            Mutiplier to apply to the quantiles.

        max_points : int, default = 2000
            If the length of X is larger than max_points, estimate the quantile on a random subset of size max_points to
            avoid memory overloads.

        Returns
        -------
            epsilon: float

        """
        means = generate_missing_data(self.seed)._nanmean(X, 0)
        X_ = X.clone()
        mask = torch.isnan(X_)
        X_[mask] = (mask * means)[mask]

        np.random.seed(self.seed)
        idx = np.random.choice(len(X_), min(max_points, len(X_)), replace=False)
        X = X_[idx]
        dists = ((X[:, None] - X) ** 2).sum(2).flatten() / 2.
        dists = dists[dists > 0]

        return generate_missing_data(self.seed)._quantile(dists, quant, 0).item() * mult

    #### Accuracy Metrics ####
    def _MAE(self, X, X_true, mask):
        """
        Mean Absolute Error (MAE) between imputed variables and ground truth. Pytorch/Numpy agnostic

        Parameters
        ----------
        X : torch.DoubleTensor or np.ndarray, shape (n, d)
            Data with imputed variables.

        X_true : torch.DoubleTensor or np.ndarray, shape (n, d)
            Ground truth.

        mask : torch.BoolTensor or np.ndarray of booleans, shape (n, d)
            Missing value mask (missing if True)

        Returns
        -------
            MAE : float

        """
        if torch.is_tensor(mask):
            mask_ = mask.bool()
            return torch.abs(X[mask_] - X_true[mask_]).sum() / mask_.sum()
        else:  # should be an ndarray
            mask_ = mask.astype(bool)
            return np.absolute(X[mask_] - X_true[mask_]).sum() / mask_.sum()

    def _RMSE(self, X, X_true, mask):
        """
        Root Mean Squared Error (MAE) between imputed variables and ground truth. Pytorch/Numpy agnostic

        Parameters
        ----------
        X : torch.DoubleTensor or np.ndarray, shape (n, d)
            Data with imputed variables.

        X_true : torch.DoubleTensor or np.ndarray, shape (n, d)
            Ground truth.

        mask : torch.BoolTensor or np.ndarray of booleans, shape (n, d)
            Missing value mask (missing if True)

        Returns
        -------
            RMSE : float

        """
        if torch.is_tensor(mask):
            mask_ = mask.bool()
            return (((X[mask_] - X_true[mask_]) ** 2).sum() / mask_.sum()).sqrt()
        else:  # should be an ndarray
            mask_ = mask.astype(bool)
            return np.sqrt(((X[mask_] - X_true[mask_]) ** 2).sum() / mask_.sum())

    ##################### MISSING DATA MECHANISMS #############################

    ##### Missing At Random ######

    def _MAR_mask(self, X, p, p_obs):
        """
        Missing at random mechanism with a logistic masking model. First, a subset of variables with *no* missing values is
        randomly selected. The remaining variables have missing values according to a logistic model with random weights,
        re-scaled so as to attain the desired proportion of missing values on those variables.

        Parameters
        ----------
        X : torch.DoubleTensor or np.ndarray, shape (n, d)
            Data for which missing values will be simulated. If a numpy array is provided,
            it will be converted to a pytorch tensor.

        p : float
            Proportion of missing values to generate for variables which will have missing values.

        p_obs : float
            Proportion of variables with *no* missing values that will be used for the logistic masking model.

        Returns
        -------
        mask : torch.BoolTensor or np.ndarray (depending on type of X)
            Mask of generated missing values (True if the value is missing).

        """

        n, d = X.shape

        to_torch = torch.is_tensor(X)  ## output a pytorch tensor, or a numpy array
        if not to_torch:
            X = torch.from_numpy(X)

        mask = torch.zeros(n, d).bool() if to_torch else np.zeros((n, d)).astype(bool)

        d_obs = max(int(p_obs * d), 1)  ## number of variables that will have no missing values (at least one variable)
        d_na = d - d_obs  ## number of variables that will have missing values

        ### Sample variables that will all be observed, and those with missing values:
        np.random.seed(self.seed)
        idxs_obs = np.random.choice(d, d_obs, replace=False)
        idxs_nas = np.array([i for i in range(d) if i not in idxs_obs])

        ### Other variables will have NA proportions that depend on those observed variables, through a logistic model
        ### The parameters of this logistic model are random.

        ### Pick coefficients so that W^Tx has unit variance (avoids shrinking)
        coeffs = generate_missing_data(self.seed)._pick_coeffs(X, idxs_obs, idxs_nas)
        ### Pick the intercepts to have a desired amount of missing values
        intercepts = generate_missing_data(self.seed)._fit_intercepts(X[:, idxs_obs], coeffs, p)

        ps = torch.sigmoid(X[:, idxs_obs].mm(coeffs) + intercepts)

        ber = torch.rand(n, d_na)
        mask[:, idxs_nas] = ber < ps

        return mask

    ##### Missing not at random ######

    def _MNAR_mask_logistic(self, X, p, p_params=.3, exclude_inputs=True):
        """
        Missing not at random mechanism with a logistic masking model. It implements two mechanisms:
        (i) Missing probabilities are selected with a logistic model, taking all variables as inputs. Hence, values that are
        inputs can also be missing.
        (ii) Variables are split into a set of intputs for a logistic model, and a set whose missing probabilities are
        determined by the logistic model. Then inputs are then masked MCAR (hence, missing values from the second set will
        depend on masked values.
        In either case, weights are random and the intercept is selected to attain the desired proportion of missing values.

        Parameters
        ----------
        X : torch.DoubleTensor or np.ndarray, shape (n, d)
            Data for which missing values will be simulated.
            If a numpy array is provided, it will be converted to a pytorch tensor.

        p : float
            Proportion of missing values to generate for variables which will have missing values.

        p_params : float
            Proportion of variables that will be used for the logistic masking model (only if exclude_inputs).

        exclude_inputs : boolean, default=True
            True: mechanism (ii) is used, False: (i)

        Returns
        -------
        mask : torch.BoolTensor or np.ndarray (depending on type of X)
            Mask of generated missing values (True if the value is missing).

        """

        n, d = X.shape

        to_torch = torch.is_tensor(X)  ## output a pytorch tensor, or a numpy array
        if not to_torch:
            X = torch.from_numpy(X)
        mask = torch.zeros(n, d).bool() if to_torch else np.zeros((n, d)).astype(bool)
        d_params = max(int(p_params * d), 1) if exclude_inputs else d  ## number of variables used as inputs (at least 1)
        d_na = d - d_params if exclude_inputs else d  ## number of variables masked with the logistic model

        ### Sample variables that will be parameters for the logistic regression:
        np.random.seed(self.seed)
        idxs_params = np.random.choice(d, d_params, replace=False) if exclude_inputs else np.arange(d)
        idxs_nas = np.array([i for i in range(d) if i not in idxs_params]) if exclude_inputs else np.arange(d)

        ### Other variables will have NA proportions selected by a logistic model
        ### The parameters of this logistic model are random.

        ### Pick coefficients so that W^Tx has unit variance (avoids shrinking)
        coeffs = generate_missing_data(self.seed)._pick_coeffs(X, idxs_params, idxs_nas)
        ### Pick the intercepts to have a desired amount of missing values
        intercepts = generate_missing_data(self.seed)._fit_intercepts(X[:, idxs_params], coeffs, p)

        ps = torch.sigmoid(X[:, idxs_params].mm(coeffs) + intercepts)

        ber = torch.rand(n, d_na)
        mask[:, idxs_nas] = ber < ps

        ## If the inputs of the logistic model are excluded from MNAR missingness,
        ## mask some values used in the logistic model at random.
        ## This makes the missingness of other variables potentially dependent on masked values

        if exclude_inputs:
            mask[:, idxs_params] = torch.rand(n, d_params) < p

        return mask

    def _MNAR_self_mask_logistic(self, X, p):
        """
        Missing not at random mechanism with a logistic self-masking model. Variables have missing values probabilities
        given by a logistic model, taking the same variable as input (hence, missingness is independent from one variable
        to another). The intercepts are selected to attain the desired missing rate.

        Parameters
        ----------
        X : torch.DoubleTensor or np.ndarray, shape (n, d)
            Data for which missing values will be simulated.
            If a numpy array is provided, it will be converted to a pytorch tensor.

        p : float
            Proportion of missing values to generate for variables which will have missing values.

        Returns
        -------
        mask : torch.BoolTensor or np.ndarray (depending on type of X)
            Mask of generated missing values (True if the value is missing).

        """

        n, d = X.shape

        to_torch = torch.is_tensor(X)  ## output a pytorch tensor, or a numpy array
        if not to_torch:
            X = torch.from_numpy(X)

        ### Variables will have NA proportions that depend on those observed variables, through a logistic model
        ### The parameters of this logistic model are random.

        ### Pick coefficients so that W^Tx has unit variance (avoids shrinking)
        coeffs = generate_missing_data(self.seed)._pick_coeffs(X, self_mask=True)
        ### Pick the intercepts to have a desired amount of missing values
        intercepts = generate_missing_data(self.seed)._fit_intercepts(X, coeffs, p, self_mask=True)

        ps = torch.sigmoid(X * coeffs + intercepts)
        np.random.seed(self.seed)
        ber = torch.rand(n, d) if to_torch else np.random.rand(n, d)
        mask = ber < ps if to_torch else ber < ps.numpy()

        return mask

    def _MNAR_mask_quantiles(self, X, p, q, p_params, cut='both', MCAR=False):
        """
        Missing not at random mechanism with quantile censorship. First, a subset of variables which will have missing
        variables is randomly selected. Then, missing values are generated on the q-quantiles at random. Since
        missingness depends on quantile information, it depends on masked values, hence this is a MNAR mechanism.

        Parameters
        ----------
        X : torch.DoubleTensor or np.ndarray, shape (n, d)
            Data for which missing values will be simulated.
            If a numpy array is provided, it will be converted to a pytorch tensor.

        p : float
            Proportion of missing values to generate for variables which will have missing values.

        q : float
            Quantile level at which the cuts should occur

        p_params : float
            Proportion of variables that will have missing values

        cut : 'both', 'upper' or 'lower', default = 'both'
            Where the cut should be applied. For instance, if q=0.25 and cut='upper', then missing values will be generated
            in the upper quartiles of selected variables.

        MCAR : bool, default = True
            If true, masks variables that were not selected for quantile censorship with a MCAR mechanism.

        Returns
        -------
        mask : torch.BoolTensor or np.ndarray (depending on type of X)
            Mask of generated missing values (True if the value is missing).

        """
        n, d = X.shape

        to_torch = torch.is_tensor(X)  ## output a pytorch tensor, or a numpy array
        if not to_torch:
            X = torch.from_numpy(X)

        mask = torch.zeros(n, d).bool() if to_torch else np.zeros((n, d)).astype(bool)

        d_na = max(int(p_params * d), 1)  ## number of variables that will have NMAR values

        ### Sample variables that will have imps at the extremes
        np.random.seed(self.seed)
        idxs_na = np.random.choice(d, d_na, replace=False)  ### select at least one variable with missing values

        ### check if values are greater/smaller that corresponding quantiles
        if cut == 'upper':
            quants = generate_missing_data(self.seed)._quantile(X[:, idxs_na], 1 - q, dim=0)
            m = X[:, idxs_na] >= quants
        elif cut == 'lower':
            quants = generate_missing_data(self.seed)._quantile(X[:, idxs_na], q, dim=0)
            m = X[:, idxs_na] <= quants
        elif cut == 'both':
            u_quants = generate_missing_data(self.seed)._quantile(X[:, idxs_na], 1 - q, dim=0)
            l_quants = generate_missing_data(self.seed)._quantile(X[:, idxs_na], q, dim=0)
            m = (X[:, idxs_na] <= l_quants) | (X[:, idxs_na] >= u_quants)

        ### Hide some values exceeding quantiles
        ber = torch.rand(n, d_na)
        mask[:, idxs_na] = (ber < p) & m

        if MCAR:
            ## Add a mcar mecanism on top
            mask = mask | (torch.rand(n, d) < p)

        return mask

    def _pick_coeffs(self, X, idxs_obs=None, idxs_nas=None, self_mask=False):
        n, d = X.shape
        if self_mask:
            coeffs = torch.randn(d)
            Wx = X * coeffs
            coeffs /= torch.std(Wx, 0)
        else:
            d_obs = len(idxs_obs)
            d_na = len(idxs_nas)
            coeffs = torch.randn(d_obs, d_na)
            Wx = X[:, idxs_obs].mm(coeffs)
            coeffs /= torch.std(Wx, 0, keepdim=True)
        return coeffs

    def _fit_intercepts(self, X, coeffs, p, self_mask=False):
        if self_mask:
            d = len(coeffs)
            intercepts = torch.zeros(d)
            for j in range(d):
                def f(x):
                    return torch.sigmoid(X * coeffs[j] + x).mean().item() - p

                intercepts[j] = optimize.bisect(f, -50, 50)
        else:
            d_obs, d_na = coeffs.shape
            intercepts = torch.zeros(d_na)
            for j in range(d_na):
                def f(x):
                    return torch.sigmoid(X.mv(coeffs[:, j]) + x).mean().item() - p

                intercepts[j] = optimize.bisect(f, -50, 50)
        return intercepts



