"""
Author: Alicia Curth
Module implements X-learner from Kuenzel et al (2019) using NNs
"""
import pandas as pd

from catenets.models.constants import *
from catenets.models.base import BaseCATENet, train_output_net_only
from catenets.models.model_utils import check_shape_1d_data, check_X_is_np


class XNet(BaseCATENet):
    """
    Class implements X-learner using NNs.

    Parameters
    ----------
    weight_strategy: int, default None
        Which strategy to use to weight the two CATE estimators in the second stage. weight_strategy
        is coded as follows: for tau(x)=g(x)tau_0(x) + (1-g(x))tau_1(x) [eq 9, kuenzel et al (2019)]
        weight_strategy=0 sets g(x)=0, weight_strategy=1 sets g(x)=1,
        weight_strategy=None sets g(x)=pi(x) [propensity score],
         weight_strategy=-1 sets g(x)=(1-pi(x))
    binary_y: bool, default False
        Whether the outcome is binary
    n_layers_out: int
        First stage Number of hypothesis layers (n_layers_out x n_units_out + 1 x Dense layer)
    n_units_out: int
        First stage Number of hidden units in each hypothesis layer
    n_layers_r: int
        First stage Number of representation layers before hypothesis layers (distinction between
        hypothesis layers and representation layers is made to match TARNet & SNets)
    n_units_r: int
        First stage Number of hidden units in each representation layer
    n_layers_out_t: int
        Second stage Number of hypothesis layers (n_layers_out x n_units_out + 1 x Dense layer)
    n_units_out_t: int
        Second stage Number of hidden units in each hypothesis layer
    n_layers_r_t: int
        Second stage Number of representation layers before hypothesis layers (distinction between
        hypothesis layers and representation layers is made to match TARNet & SNets)
    n_units_r_t: int
        Second stage Number of hidden units in each representation layer
    penalty_l2: float
        First stage l2 (ridge) penalty
    penalty_l2_t: float
        Second stage l2 (ridge) penalty
    step_size: float
        First stage learning rate for optimizer
    step_size_t: float
        Second stage learning rate for optimizer
    n_iter: int
        Maximum number of iterations
    batch_size: int
        Batch size
    val_split_prop: float
        Proportion of samples used for validation split (can be 0)
    early_stopping: bool, default True
        Whether to use early stopping
    patience: int
        Number of iterations to wait before early stopping after decrease in validation loss
    n_iter_min: int
        Minimum number of iterations to go through before starting early stopping
    verbose: int, default 1
        Whether to print notifications
    n_iter_print: int
        Number of iterations after which to print updates
    seed: int
        Seed used
    nonlin: string, default 'elu'
        Nonlinearity to use in NN
    """
    def __init__(self, weight_strategy: int = None,
                 binary_y: bool = False,
                 n_layers_out: int = DEFAULT_LAYERS_OUT,
                 n_layers_r: int = DEFAULT_LAYERS_R,
                 n_layers_out_t: int = DEFAULT_LAYERS_OUT_T,
                 n_layers_r_t: int = DEFAULT_LAYERS_R_T,
                 n_units_out: int = DEFAULT_UNITS_OUT,
                 n_units_r: int = DEFAULT_UNITS_R,
                 n_units_out_t: int = DEFAULT_UNITS_OUT_T,
                 n_units_r_t: int = DEFAULT_UNITS_R_T,
                 penalty_l2: float = DEFAULT_PENALTY_L2,
                 penalty_l2_t: float = DEFAULT_PENALTY_L2,
                 step_size: float = DEFAULT_STEP_SIZE,
                 step_size_t: float = DEFAULT_STEP_SIZE_T,
                 n_iter: int = DEFAULT_N_ITER,
                 batch_size: int = DEFAULT_BATCH_SIZE,
                 n_iter_min: int = DEFAULT_N_ITER_MIN,
                 val_split_prop: float = DEFAULT_VAL_SPLIT,
                 early_stopping: bool = True,
                 patience: int = DEFAULT_PATIENCE,
                 verbose: int = 1, n_iter_print: int = DEFAULT_N_ITER_PRINT,
                 seed: int = DEFAULT_SEED, nonlin: str = DEFAULT_NONLIN):
        # settings
        self.weight_strategy = weight_strategy
        self.binary_y = binary_y

        # model architecture hyperparams
        self.n_layers_out = n_layers_out
        self.n_layers_out_t = n_layers_out_t
        self.n_layers_r = n_layers_r
        self.n_layers_r_t = n_layers_r_t
        self.n_units_out = n_units_out
        self.n_units_out_t = n_units_out_t
        self.n_units_r = n_units_r
        self.n_units_r_t = n_units_r_t
        self.nonlin = nonlin

        # other hyperparameters
        self.penalty_l2 = penalty_l2
        self.penalty_l2_t = penalty_l2_t
        self.step_size = step_size
        self.step_size_t = step_size_t
        self.n_iter = n_iter
        self.batch_size = batch_size
        self.verbose = verbose
        self.n_iter_print = n_iter_print
        self.seed = seed
        self.val_split_prop = val_split_prop
        self.early_stopping = early_stopping
        self.patience = patience
        self.n_iter_min = n_iter_min

    def _get_train_function(self):
        return train_x_net

    def _get_predict_function(self):
        # Two step nets do not need this
        return predict_x_net

    def predict(self, X, return_po: bool = False, return_prop: bool = False):
        """
        Predict treatment effect estimates using a CATENet. Depending on method, can also return
        potential outcome estimate and propensity score estimate.

        Parameters
        ----------
        X: pd.DataFrame or np.array
            Covariate matrix
        return_po: bool, default False
            Whether to return potential outcome estimate
        return_prop: bool, default False
            Whether to return propensity estimate

        Returns
        -------
        array of CATE estimates, optionally also potential outcomes and propensity
        """
        X = check_X_is_np(X)
        predict_func = self._get_predict_function()
        return predict_func(X, trained_params=self._params,
                            predict_funs=self._predict_funs,
                            return_po=return_po, return_prop=return_prop,
                            weight_strategy=self.weight_strategy)


def train_x_net(X, y, w, weight_strategy: int = None,
                binary_y: bool = False,
                n_layers_out: int = DEFAULT_LAYERS_OUT,
                n_layers_r: int = DEFAULT_LAYERS_R,
                n_layers_out_t: int = DEFAULT_LAYERS_OUT_T,
                n_layers_r_t: int = DEFAULT_LAYERS_R_T,
                n_units_out: int = DEFAULT_UNITS_OUT,
                n_units_r: int = DEFAULT_UNITS_R,
                n_units_out_t: int = DEFAULT_UNITS_OUT_T,
                n_units_r_t: int = DEFAULT_UNITS_R_T,
                penalty_l2: float = DEFAULT_PENALTY_L2,
                penalty_l2_t: float = DEFAULT_PENALTY_L2,
                step_size: float = DEFAULT_STEP_SIZE,
                step_size_t: float = DEFAULT_STEP_SIZE_T,
                n_iter: int = DEFAULT_N_ITER,
                batch_size: int = DEFAULT_BATCH_SIZE,
                n_iter_min: int = DEFAULT_N_ITER_MIN,
                val_split_prop: float = DEFAULT_VAL_SPLIT,
                early_stopping: bool = True,
                patience: int = DEFAULT_PATIENCE,
                verbose: int = 1, n_iter_print: int = DEFAULT_N_ITER_PRINT,
                seed: int = DEFAULT_SEED, nonlin: str = DEFAULT_NONLIN,
                return_val_loss: bool = False, avg_objective: bool = DEFAULT_AVG_OBJECTIVE):
    y = check_shape_1d_data(y)
    if len(w.shape) > 1:
        w = w.reshape((len(w),))

    if weight_strategy not in [0, 1, -1, None]:
        # weight_strategy is coded as follows:
        # for tau(x)=g(x)tau_0(x) + (1-g(x))tau_1(x) [eq 9, kuenzel et al (2019)]
        # weight_strategy=0 sets g(x)=0, weight_strategy=1 sets g(x)=1,
        # weight_strategy=None sets g(x)=pi(x) [propensity score],
        # weight_strategy=-1 sets g(x)=(1-pi(x))
        raise ValueError('XNet only implements weight_strategy in [0, 1, -1, None]')

    # first stage: get estimates of PO regression
    if verbose > 0:
        print("Training first stage")

    if not weight_strategy == 1:
        if verbose > 0:
            print('Training PO_0 Net')
        params_0, predict_fun_0 = train_output_net_only(X[w == 0], y[w == 0],
                                                        binary_y=binary_y,
                                                        n_layers_out=n_layers_out,
                                                        n_units_out=n_units_out,
                                                        n_layers_r=n_layers_r,
                                                        n_units_r=n_units_r,
                                                        penalty_l2=penalty_l2,
                                                        step_size=step_size,
                                                        n_iter=n_iter,
                                                        batch_size=batch_size,
                                                        val_split_prop=val_split_prop,
                                                        early_stopping=early_stopping,
                                                        patience=patience,
                                                        n_iter_min=n_iter_min,
                                                        n_iter_print=n_iter_print,
                                                        verbose=verbose,
                                                        seed=seed, nonlin=nonlin,
                                                        avg_objective=avg_objective)
        mu_hat_0 = predict_fun_0(params_0, X[w == 1])
    else:
        mu_hat_0 = None

    if not weight_strategy == 0:
        if verbose > 0:
            print('Training PO_1 Net')
        params_1, predict_fun_1 = train_output_net_only(X[w == 1], y[w == 1],
                                                        binary_y=binary_y,
                                                        n_layers_out=n_layers_out,
                                                        n_units_out=n_units_out,
                                                        n_layers_r=n_layers_r,
                                                        n_units_r=n_units_r,
                                                        penalty_l2=penalty_l2,
                                                        step_size=step_size,
                                                        n_iter=n_iter,
                                                        batch_size=batch_size,
                                                        val_split_prop=val_split_prop,
                                                        early_stopping=early_stopping,
                                                        patience=patience,
                                                        n_iter_min=n_iter_min,
                                                        n_iter_print=n_iter_print,
                                                        verbose=verbose,
                                                        seed=seed, nonlin=nonlin,
                                                        avg_objective=avg_objective)
        mu_hat_1 = predict_fun_1(params_1, X[w == 0])
    else:
        mu_hat_1 = None

    if weight_strategy is None or weight_strategy == -1:
        # also fit propensity estimator
        if verbose > 0:
            print('Training propensity net')
        params_prop, predict_fun_prop = train_output_net_only(X, w,
                                                              binary_y=True,
                                                              n_layers_out=n_layers_out,
                                                              n_units_out=n_units_out,
                                                              n_layers_r=n_layers_r,
                                                              n_units_r=n_units_r,
                                                              penalty_l2=penalty_l2,
                                                              step_size=step_size,
                                                              n_iter=n_iter,
                                                              batch_size=batch_size,
                                                              val_split_prop=val_split_prop,
                                                              early_stopping=early_stopping,
                                                              patience=patience,
                                                              n_iter_min=n_iter_min,
                                                              n_iter_print=n_iter_print,
                                                              verbose=verbose,
                                                              seed=seed, nonlin=nonlin,
                                                              avg_objective=avg_objective)

    else:
        params_prop, predict_fun_prop = None, None

    # second stage
    if verbose > 0:
        print("Training second stage")
    if not weight_strategy == 0:
        # fit tau_0
        if verbose > 0:
            print("Fitting tau_0")
        pseudo_outcome0 = mu_hat_1 - y[w == 0]
        params_tau0, predict_fun_tau0 = train_output_net_only(X[w == 0], pseudo_outcome0,
                                                              binary_y=False,
                                                              n_layers_out=n_layers_out_t,
                                                              n_units_out=n_units_out_t,
                                                              n_layers_r=n_layers_r_t,
                                                              n_units_r=n_units_r_t,
                                                              penalty_l2=penalty_l2_t,
                                                              step_size=step_size_t,
                                                              n_iter=n_iter,
                                                              batch_size=batch_size,
                                                              val_split_prop=val_split_prop,
                                                              early_stopping=early_stopping,
                                                              patience=patience,
                                                              n_iter_min=n_iter_min,
                                                              n_iter_print=n_iter_print,
                                                              verbose=verbose,
                                                              seed=seed,
                                                              return_val_loss=return_val_loss,
                                                              nonlin=nonlin,
                                                              avg_objective=avg_objective)
    else:
        params_tau0, predict_fun_tau0 = None, None

    if not weight_strategy == 1:
        # fit tau_1
        if verbose > 0:
            print("Fitting tau_1")
        pseudo_outcome1 = y[w == 1] - mu_hat_0
        params_tau1, predict_fun_tau1 = train_output_net_only(X[w == 1], pseudo_outcome1,
                                                              binary_y=False,
                                                              n_layers_out=n_layers_out_t,
                                                              n_units_out=n_units_out_t,
                                                              n_layers_r=n_layers_r_t,
                                                              n_units_r=n_units_r_t,
                                                              penalty_l2=penalty_l2_t,
                                                              step_size=step_size_t,
                                                              n_iter=n_iter,
                                                              batch_size=batch_size,
                                                              val_split_prop=val_split_prop,
                                                              early_stopping=early_stopping,
                                                              patience=patience,
                                                              n_iter_min=n_iter_min,
                                                              n_iter_print=n_iter_print,
                                                              verbose=verbose,
                                                              seed=seed,
                                                              return_val_loss=return_val_loss,
                                                              nonlin=nonlin,
                                                              avg_objective=avg_objective)

    else:
        params_tau1, predict_fun_tau1 = None, None

    params = params_tau0, params_tau1, params_prop
    predict_funs = predict_fun_tau0, predict_fun_tau1, predict_fun_prop

    return params, predict_funs


def predict_x_net(X, trained_params, predict_funs,
                  return_po: bool = False, return_prop: bool = False,
                  weight_strategy: int = None):
    if return_po:
        raise NotImplementedError('TwoStepNets have no Potential outcome predictors.')

    if return_prop:
        raise NotImplementedError('TwoStepNets have no Propensity predictors.')

    params_tau0, params_tau1, params_prop = trained_params
    predict_fun_tau0, predict_fun_tau1, predict_fun_prop = predict_funs

    if not weight_strategy == 0:
        tau0_pred = predict_fun_tau0(params_tau0, X)
    else:
        tau0_pred = 0

    if not weight_strategy == 1:
        tau1_pred = predict_fun_tau1(params_tau1, X)
    else:
        tau1_pred = 0

    if weight_strategy is None or weight_strategy == -1:
        prop_pred = predict_fun_prop(params_prop, X)

    if weight_strategy is None:
        weight = prop_pred
    elif weight_strategy == -1:
        weight = 1 - prop_pred
    elif weight_strategy == 0:
        weight = 0
    elif weight_strategy == 1:
        weight = 1

    return weight * tau0_pred + (1-weight) * tau1_pred
