import unittest
import pytest
import pickle
from sklearn.linear_model import LinearRegression, Lasso, LassoCV, LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, FunctionTransformer, PolynomialFeatures
from sklearn.model_selection import KFold, GroupKFold
from econml.dml import DML, LinearDML, SparseLinearDML, KernelDML
from econml.dml import NonParamDML, ForestDML
from econml.drlearner import DRLearner, SparseLinearDRLearner, LinearDRLearner, ForestDRLearner
from econml.ortho_iv import DMLATEIV, ProjectedDMLATEIV, DMLIV, NonParamDMLIV,\
    IntentToTreatDRIV, LinearIntentToTreatDRIV
import numpy as np
from econml.utilities import shape, hstack, vstack, reshape, cross_product
from econml.inference import BootstrapInference
from contextlib import ExitStack
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
import itertools
from econml.sklearn_extensions.linear_model import WeightedLasso, StatsModelsRLM
from econml.tests.test_statsmodels import _summarize
import econml.tests.utilities  # bugfix for assertWarns
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier


class TestRandomState(unittest.TestCase):

    @staticmethod
    def _make_data(n, p):
        np.random.seed(1283)
        X = np.random.uniform(-1, 1, size=(n, p))
        W = np.random.uniform(-1, 1, size=(n, p))

        def true_propensity(x):
            return .4 + .1 * (x[:, 0] > 0)

        def true_effect(x):
            return .4 + .2 * x[:, 0]

        def true_conf(x):
            return x[:, 1]

        T = np.random.binomial(1, true_propensity(X))
        Y = true_effect(X) * T + true_conf(X) + np.random.normal(size=(n,))
        X_test = np.zeros((100, p))
        X_test[:, 0] = np.linspace(-1, 1, 100)
        return Y, T, X, W, X_test

    @staticmethod
    def _test_random_state(est, X_test, Y, T, **kwargs):
        est.fit(Y, T, **kwargs)
        te1 = est.effect(X_test)
        est.fit(Y, T, **kwargs)
        te2 = est.effect(X_test)
        est.fit(Y, T, **kwargs)
        te3 = est.effect(X_test)
        np.testing.assert_array_equal(te1, te2, err_msg='random state fixing does not work')
        np.testing.assert_array_equal(te1, te3, err_msg='random state fixing does not work')

    def test_dml_random_state(self):
        Y, T, X, W, X_test = TestRandomState._make_data(500, 2)
        for est in [
                NonParamDML(model_y=RandomForestRegressor(n_estimators=10, max_depth=4, random_state=123),
                            model_t=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                            model_final=RandomForestRegressor(max_depth=3, n_estimators=10, min_samples_leaf=100,
                                                              bootstrap=True, random_state=123),
                            discrete_treatment=True, n_splits=2, random_state=123),
                ForestDML(model_y=RandomForestRegressor(n_estimators=10, max_depth=4, random_state=123),
                          model_t=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                          n_estimators=10,
                          discrete_treatment=True, n_crossfit_splits=2, random_state=123),
                LinearDML(model_y=RandomForestRegressor(n_estimators=10, max_depth=4, random_state=123),
                          model_t=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                          discrete_treatment=True, n_splits=2, random_state=123),
                SparseLinearDML(discrete_treatment=True, n_splits=2, random_state=123),
                KernelDML(discrete_treatment=True, n_splits=2, random_state=123)]:
            TestRandomState._test_random_state(est, X_test, Y, T, X=X, W=W)

    def test_dr_random_state(self):
        Y, T, X, W, X_test = self._make_data(500, 2)
        for est in [
                DRLearner(model_final=RandomForestRegressor(max_depth=3, n_estimators=10, min_samples_leaf=100,
                                                            bootstrap=True, random_state=123),
                          n_splits=2, random_state=123),
                LinearDRLearner(random_state=123),
                SparseLinearDRLearner(n_splits=2, random_state=123),
                ForestDRLearner(model_regression=RandomForestRegressor(n_estimators=10, max_depth=4,
                                                                       random_state=123),
                                model_propensity=RandomForestClassifier(
                    n_estimators=10, max_depth=4, random_state=123),
                    n_crossfit_splits=2, random_state=123)]:
            TestRandomState._test_random_state(est, X_test, Y, T, X=X, W=W)

    def test_orthoiv_random_state(self):
        Y, T, X, W, X_test = self._make_data(500, 2)
        for est in [
            DMLATEIV(model_Y_W=RandomForestRegressor(n_estimators=10, max_depth=4, random_state=123),
                     model_T_W=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                     model_Z_W=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                     discrete_treatment=True, discrete_instrument=True, n_splits=2, random_state=123),
            ProjectedDMLATEIV(model_Y_W=RandomForestRegressor(n_estimators=10, max_depth=4, random_state=123),
                              model_T_W=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                              model_T_WZ=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                              discrete_treatment=True, discrete_instrument=True, n_splits=2, random_state=123)]:
            TestRandomState._test_random_state(est, None, Y, T, W=W, Z=T)
        for est in [
                DMLIV(model_Y_X=RandomForestRegressor(n_estimators=10, max_depth=4, random_state=123),
                      model_T_X=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                      model_T_XZ=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                      model_final=LinearRegression(fit_intercept=False),
                      discrete_treatment=True, discrete_instrument=True, n_splits=2, random_state=123),
                NonParamDMLIV(model_Y_X=RandomForestRegressor(n_estimators=10, max_depth=4, random_state=123),
                              model_T_X=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                              model_T_XZ=RandomForestClassifier(n_estimators=10, max_depth=4, random_state=123),
                              model_final=LinearRegression(),
                              discrete_treatment=True, discrete_instrument=True, n_splits=2, random_state=123)]:
            TestRandomState._test_random_state(est, X_test, Y, T, X=X, Z=T)
        for est in [IntentToTreatDRIV(model_Y_X=RandomForestRegressor(n_estimators=10, max_depth=4, random_state=123),
                                      model_T_XZ=RandomForestClassifier(n_estimators=10,
                                                                        max_depth=4, random_state=123),
                                      flexible_model_effect=RandomForestRegressor(n_estimators=10,
                                                                                  max_depth=4, random_state=123),
                                      n_splits=2, random_state=123),
                    LinearIntentToTreatDRIV(model_Y_X=RandomForestRegressor(n_estimators=10,
                                                                            max_depth=4, random_state=123),
                                            model_T_XZ=RandomForestClassifier(n_estimators=10,
                                                                              max_depth=4, random_state=123),
                                            flexible_model_effect=RandomForestRegressor(n_estimators=10,
                                                                                        max_depth=4,
                                                                                        random_state=123),
                                            n_splits=2, random_state=123)]:
            TestRandomState._test_random_state(est, X_test, Y, T, X=X, W=W, Z=T)
