# Copyright (c) PyWhy contributors. All rights reserved.
# Licensed under the MIT License.
import abc
import numpy as np
from econml.utilities import cross_product
from statsmodels.tools.tools import add_constant
try:
    import matplotlib
    import matplotlib.pyplot as plt
except ImportError as exn:
    from .utilities import MissingModule

    # make any access to matplotlib or plt throw an exception
    matplotlib = plt = MissingModule("matplotlib is no longer a dependency of the main econml package; "
                                     "install econml[plt] or econml[all] to require it, or install matplotlib "
                                     "separately, to use the tree interpreters", exn)


class _BaseDynamicPanelDGP:

    def __init__(self, n_periods, n_treatments, n_x):
        self.n_periods = n_periods
        self.n_treatments = n_treatments
        self.n_x = n_x
        return

    @abc.abstractmethod
    def create_instance(self, *args, **kwargs):
        pass

    @abc.abstractmethod
    def _gen_data_with_policy(self, n_units, policy_gen, random_seed=123):
        pass

    def static_policy_data(self, n_units, tau, random_seed=123):
        def policy_gen(Tpre, X, period):
            return tau[period]
        return self._gen_data_with_policy(n_units, policy_gen, random_seed=random_seed)

    def adaptive_policy_data(self, n_units, policy_gen, random_seed=123):
        return self._gen_data_with_policy(n_units, policy_gen, random_seed=random_seed)

    def static_policy_effect(self, tau, mc_samples=1000):
        Y_tau, _, _, _ = self.static_policy_data(mc_samples, tau)
        Y_zero, _, _, _ = self.static_policy_data(
            mc_samples, np.zeros((self.n_periods, self.n_treatments)))
        return np.mean(Y_tau[np.arange(Y_tau.shape[0]) % self.n_periods == self.n_periods - 1]) - \
            np.mean(Y_zero[np.arange(Y_zero.shape[0]) %
                           self.n_periods == self.n_periods - 1])

    def adaptive_policy_effect(self, policy_gen, mc_samples=1000):
        Y_tau, _, _, _ = self.adaptive_policy_data(mc_samples, policy_gen)
        Y_zero, _, _, _ = self.static_policy_data(
            mc_samples, np.zeros((self.n_periods, self.n_treatments)))
        return np.mean(Y_tau[np.arange(Y_tau.shape[0]) % self.n_periods == self.n_periods - 1]) - \
            np.mean(Y_zero[np.arange(Y_zero.shape[0]) %
                           self.n_periods == self.n_periods - 1])


class DynamicPanelDGP(_BaseDynamicPanelDGP):

    def __init__(self, n_periods, n_treatments, n_x):
        super().__init__(n_periods, n_treatments, n_x)

    def create_instance(self, s_x, sigma_x=.8, sigma_y=.1, conf_str=5, hetero_strength=.5, hetero_inds=None,
                        autoreg=.25, state_effect=.25, random_seed=123):
        np.random.seed(random_seed)
        self.s_x = s_x
        self.conf_str = conf_str
        self.sigma_x = sigma_x
        self.sigma_y = sigma_y
        self.hetero_inds = hetero_inds.astype(
            int) if hetero_inds is not None else hetero_inds
        self.endo_inds = np.setdiff1d(
            np.arange(self.n_x), hetero_inds).astype(int)
        # The first s_x state variables are confounders. The final s_x variables are exogenous and can create
        # heterogeneity
        self.Alpha = np.random.uniform(-1, 1,
                                       size=(self.n_x, self.n_treatments))
        self.Alpha /= np.linalg.norm(self.Alpha, axis=1, ord=1, keepdims=True)
        self.Alpha *= state_effect
        if self.hetero_inds is not None:
            self.Alpha[self.hetero_inds] = 0

        self.Beta = np.zeros((self.n_x, self.n_x))
        for t in range(self.n_x):
            self.Beta[t, :] = autoreg * np.roll(np.random.uniform(low=4.0**(-np.arange(
                0, self.n_x)), high=4.0**(-np.arange(1, self.n_x + 1))), t)
        if self.hetero_inds is not None:
            self.Beta[np.ix_(self.endo_inds, self.hetero_inds)] = 0
            self.Beta[np.ix_(self.hetero_inds, self.endo_inds)] = 0

        self.epsilon = np.random.uniform(-1, 1, size=self.n_treatments)
        self.zeta = np.zeros(self.n_x)
        self.zeta[:self.s_x] = self.conf_str / self.s_x

        self.y_hetero_effect = np.zeros(self.n_x)
        self.x_hetero_effect = np.zeros(self.n_x)
        if self.hetero_inds is not None:
            self.y_hetero_effect[self.hetero_inds] = np.random.uniform(.5 * hetero_strength, 1.5 * hetero_strength) / \
                len(self.hetero_inds)
            self.x_hetero_effect[self.hetero_inds] = np.random.uniform(.5 * hetero_strength, 1.5 * hetero_strength) / \
                len(self.hetero_inds)

        self.true_effect = np.zeros((self.n_periods, self.n_treatments))
        # Invert indices to match latest API
        self.true_effect[self.n_periods - 1] = self.epsilon
        for t in np.arange(self.n_periods - 2, -1, -1):
            self.true_effect[t, :] = (self.zeta.reshape(
                1, -1) @ np.linalg.matrix_power(self.Beta, (self.n_periods - 1 - t) - 1) @ self.Alpha)

        self.true_hetero_effect = np.zeros(
            (self.n_periods, (self.n_x + 1) * self.n_treatments))
        self.true_hetero_effect[self.n_periods - 1, :] = cross_product(
            add_constant(self.y_hetero_effect.reshape(1, -1), has_constant='add'),
            self.epsilon.reshape(1, -1))
        for t in np.arange(self.n_periods - 2, -1, -1):
            # Invert indices to match latest API
            self.true_hetero_effect[t, :] = cross_product(
                add_constant(self.x_hetero_effect.reshape(1, -1), has_constant='add'),
                self.zeta.reshape(1, -1) @ np.linalg.matrix_power(
                    self.Beta, (self.n_periods - 1 - t) - 1) @ self.Alpha)
        return self

    def hetero_effect_fn(self, t, x):
        if t == self.n_periods - 1:
            return (np.dot(self.y_hetero_effect, x.flatten()) + 1) * self.epsilon
        else:
            return (np.dot(self.x_hetero_effect, x.flatten()) + 1) *\
                (self.zeta.reshape(1, -1) @ np.linalg.matrix_power(self.Beta, (self.n_periods - 1 - t) - 1)
                    @ self.Alpha).flatten()

    def _gen_data_with_policy(self, n_units, policy_gen, random_seed=123):
        np.random.seed(random_seed)
        Y = np.zeros(n_units * self.n_periods)
        T = np.zeros((n_units * self.n_periods, self.n_treatments))
        X = np.zeros((n_units * self.n_periods, self.n_x))
        groups = np.zeros(n_units * self.n_periods)
        for t in range(n_units * self.n_periods):
            period = t % self.n_periods
            if period == 0:
                X[t] = np.random.normal(0, self.sigma_x, size=self.n_x)
                const_x0 = X[t][self.hetero_inds]
                T[t] = policy_gen(np.zeros(self.n_treatments), X[t], period)
            else:
                X[t] = (np.dot(self.x_hetero_effect, X[t - 1]) + 1) * np.dot(self.Alpha, T[t - 1]) + \
                    np.dot(self.Beta, X[t - 1]) + \
                    np.random.normal(0, self.sigma_x, size=self.n_x)
                # The feature for heterogeneity stays constant
                X_t = X[t].copy()
                X_t[self.hetero_inds] = const_x0
                T[t] = policy_gen(T[t - 1], X[t], period)
            Y[t] = (np.dot(self.y_hetero_effect, X_t if period != 0 else X[t]) + 1) * np.dot(self.epsilon, T[t]) + \
                np.dot(X[t], self.zeta) + \
                np.random.normal(0, self.sigma_y)
            groups[t] = t // self.n_periods

        return Y, T, X[:, self.hetero_inds] if (self.hetero_inds is not None) else None, X[:, self.endo_inds], groups

    def observational_data(self, n_units, gamma=0, s_t=1, sigma_t=0.5, random_seed=123):
        """Generate observational data with some observational treatment policy parameters.

        Parameters
        ----------
        n_units : how many units to observe
        gamma : what is the degree of auto-correlation of the treatments across periods
        s_t : sparsity of treatment policy; how many states does it depend on
        sigma_t : what is the std of the exploration/randomness in the treatment
        """
        Delta = np.zeros((self.n_treatments, self.n_x))
        Delta[:, :s_t] = self.conf_str / s_t

        def policy_gen(Tpre, X, period):
            return gamma * Tpre + (1 - gamma) * np.dot(Delta, X) + \
                np.random.normal(0, sigma_t, size=self.n_treatments)
        return self._gen_data_with_policy(n_units, policy_gen, random_seed=random_seed)


# Auxiliary function for adding xticks and vertical lines when plotting results
# for dynamic dml vs ground truth parameters.
def add_vlines(n_periods, n_treatments, hetero_inds):
    locs, labels = plt.xticks([], [])
    locs += [- .501 + (len(hetero_inds) + 1) / 2]
    labels += ["\n\n$\\tau_{{{}}}$".format(0)]
    locs += [qx for qx in np.arange(len(hetero_inds) + 1)]
    labels += ["$1$"] + ["$x_{{{}}}$".format(qx) for qx in hetero_inds]
    for q in np.arange(1, n_treatments):
        plt.axvline(x=q * (len(hetero_inds) + 1) - .5,
                    linestyle='--', color='red', alpha=.2)
        locs += [q * (len(hetero_inds) + 1) - .501 + (len(hetero_inds) + 1) / 2]
        labels += ["\n\n$\\tau_{{{}}}$".format(q)]
        locs += [(q * (len(hetero_inds) + 1) + qx)
                 for qx in np.arange(len(hetero_inds) + 1)]
        labels += ["$1$"] + ["$x_{{{}}}$".format(qx) for qx in hetero_inds]
    locs += [- .501 + (len(hetero_inds) + 1) * n_treatments / 2]
    labels += ["\n\n\n\n$\\theta_{{{}}}$".format(0)]
    for t in np.arange(1, n_periods):
        plt.axvline(x=t * (len(hetero_inds) + 1) *
                    n_treatments - .5, linestyle='-', alpha=.6)
        locs += [t * (len(hetero_inds) + 1) * n_treatments - .501 +
                 (len(hetero_inds) + 1) * n_treatments / 2]
        labels += ["\n\n\n\n$\\theta_{{{}}}$".format(t)]
        locs += [t * (len(hetero_inds) + 1) *
                 n_treatments - .501 + (len(hetero_inds) + 1) / 2]
        labels += ["\n\n$\\tau_{{{}}}$".format(0)]
        locs += [t * (len(hetero_inds) + 1) * n_treatments +
                 qx for qx in np.arange(len(hetero_inds) + 1)]
        labels += ["$1$"] + ["$x_{{{}}}$".format(qx) for qx in hetero_inds]
        for q in np.arange(1, n_treatments):
            plt.axvline(x=t * (len(hetero_inds) + 1) * n_treatments + q * (len(hetero_inds) + 1) - .5,
                        linestyle='--', color='red', alpha=.2)
            locs += [t * (len(hetero_inds) + 1) * n_treatments + q *
                     (len(hetero_inds) + 1) - .501 + (len(hetero_inds) + 1) / 2]
            labels += ["\n\n$\\tau_{{{}}}$".format(q)]
            locs += [t * (len(hetero_inds) + 1) * n_treatments + (q * (len(hetero_inds) + 1) + qx)
                     for qx in np.arange(len(hetero_inds) + 1)]
            labels += ["$1$"] + ["$x_{{{}}}$".format(qx) for qx in hetero_inds]
    plt.xticks(locs, labels)
    plt.tight_layout()
