# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

"""Tests for two stage least squares module."""

import unittest
import numpy as np
import warnings
import pytest

from econml.utilities import shape, reshape
from econml.two_stage_least_squares import (NonparametricTwoStageLeastSquares, HermiteFeatures, DPolynomialFeatures)
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures


class Test2SLS(unittest.TestCase):

    def test_hermite_shape(self):
        for d, s in [(3, 0), (4, 2)]:
            for j in [True, False]:
                for n, x in [(5, 1), (7, 3)]:
                    last_dim = (d + 1)**x if j else (d + 1) * x
                    correct_shape = (n,) + (x,) * s + (last_dim,)
                    output_shape = shape(HermiteFeatures(d, s, j).fit_transform(np.zeros((n, x))))
                    assert output_shape == correct_shape

    def test_hermite_results(self):
        inputs = np.random.normal(size=(5, 1))
        hf = HermiteFeatures(3).fit_transform(inputs)
        # first polynomials are 1, x, x*x-1, x*x*x-3*x
        ones = np.ones(shape(inputs))
        polys = np.hstack([ones, inputs, inputs * inputs - ones, inputs * inputs * inputs - 3 * inputs])
        assert(np.allclose(hf, polys * np.exp(-inputs * inputs / 2)))

        for j in [True, False]:
            hf = HermiteFeatures(1, shift=1, joint=j).fit_transform(inputs)
            # first derivatives are -x, -x^2+1 (since there's just one column, joint-ness doesn't matter)
            polys = np.hstack([-inputs, -inputs * inputs + ones])
            assert(np.allclose(hf, reshape(polys * np.exp(-inputs * inputs / 2), (5, 1, 2))))

    @pytest.mark.slow
    def test_hermite_approx(self):
        n = 50000
        x = np.random.uniform(low=-0.5, high=1.5, size=(n, 2))
        y = (x[:, 1] > x[:, 0]) * \
            (0 <= x[:, 0]) * \
            (x[:, 0] <= 1) * \
            (0 <= x[:, 1]) * \
            (x[:, 1] <= 1) + \
            np.random.normal(0, 0.01, (n,))

        def err(k, j):
            hf = HermiteFeatures(k, joint=j)
            m = LinearRegression()
            m.fit(hf.fit_transform(x[:n // 2, :]), y[:n // 2])
            return ((y[n // 2:] - m.predict(hf.fit_transform(x[n // 2:, :]))) ** 2).mean()
        # TODO: test something rather than just print...
        print([(k, j, err(k, j)) for k in range(2, 15) for j in [False, True]])

    def test_2sls_shape(self):
        n = 100

        def make_random(d):
            sz = (n, d) if d >= 0 else (n,)
            return np.random.normal(size=sz)

        for d_t in [-1, 1, 2]:
            n_t = d_t if d_t > 0 else 1
            for d_y in [-1, 1, 2]:
                for d_x in [1, 5]:
                    for d_z in [1, 2]:
                        d_w = 1
                        if d_z >= n_t:
                            T, Y, X, Z, W = [make_random(d) for d in [d_t, d_y, d_x, d_z, d_w]]
                            est = NonparametricTwoStageLeastSquares(
                                t_featurizer=PolynomialFeatures(),
                                x_featurizer=PolynomialFeatures(),
                                z_featurizer=PolynomialFeatures(),
                                dt_featurizer=DPolynomialFeatures())

                            est.fit(Y, T, X=X, W=W, Z=Z)

                            eff = est.effect(X)
                            marg_eff = est.marginal_effect(T, X)

                            effect_shape = (n,) + ((d_y,) if d_y > 0 else ())
                            marginal_effect_shape = ((n if d_x else 1,) +
                                                     ((d_y,) if d_y > 0 else ()) +
                                                     ((d_t,) if d_t > 0 else()))

                            self.assertEqual(shape(marg_eff), marginal_effect_shape)
                            self.assertEqual(shape(eff), effect_shape)

    def test_marg_eff(self):
        X = np.random.normal(size=(5000, 2))
        Z = np.random.normal(size=(5000, 2))
        W = np.random.normal(size=(5000, 1))
        # Note: no noise, just testing that we can exactly recover when we ought to be able to
        T = np.hstack([np.cross(X, Z).reshape(-1, 1) + W, (np.prod(X, axis=1) + np.prod(Z, axis=1)).reshape(-1, 1)])
        Y = X * T + X**2

        est = NonparametricTwoStageLeastSquares(
            t_featurizer=PolynomialFeatures(degree=2, interaction_only=False, include_bias=True),
            x_featurizer=PolynomialFeatures(degree=2, interaction_only=False, include_bias=True),
            z_featurizer=PolynomialFeatures(degree=2, interaction_only=False, include_bias=True),
            dt_featurizer=DPolynomialFeatures(degree=2, interaction_only=False, include_bias=True))

        est.fit(Y, T, X=X, W=W, Z=Z)

        # pick some arbitrary X
        X_test = np.array([[0.3, 0.7],
                           [0.2, 0.1]])
        eff = est.effect(X_test)  # effect = (X * 1 + X^2) - (X * 0 + X^2) = X
        np.testing.assert_almost_equal(eff, X_test)

        # pick some arbitrary T
        T_test = np.array([[-0.3, 0.1],
                           [0.6, -1.2]])
        marg_eff = est.marginal_effect(T_test, X_test)  # marg effect_{i,j} = X_i if i=j, 0 otherwise
        marg_eff_truth = np.zeros((X_test.shape[0], Y.shape[1], T.shape[1]))
        marg_eff_truth[:, range(X.shape[1]), range(X.shape[1])] = X_test[:, :]
        np.testing.assert_almost_equal(marg_eff, marg_eff_truth)

    # TODO: this tests that we can run the method; how do we test that the results are reasonable?
    def test_2sls(self):
        n = 50000
        d_w = 2
        d_z = 1
        d_x = 1
        d_t = 1
        d_y = 1
        e = np.random.uniform(low=-0.5, high=0.5, size=(n, d_x))
        z = np.random.uniform(size=(n, 1))
        w = np.random.uniform(size=(n, d_w))
        a = np.random.normal(size=(d_w, d_t))
        b = np.random.normal(size=(d_w, d_y))
        x = np.random.uniform(size=(n, d_x)) + e
        p = x + z * e + w @ a + np.random.uniform(size=(n, d_t))
        y = p * x + e + w @ b

        losses = []
        marg_effs = []

        z_fresh = np.random.uniform(size=(n, d_z))
        e_fresh = np.random.uniform(low=-0.5, high=0.5, size=(n, d_x))
        x_fresh = np.random.uniform(size=(n, d_x)) + e_fresh
        w_fresh = np.random.uniform(size=(n, d_w))
        p_fresh = x_fresh + z_fresh * e_fresh + np.random.uniform(size=(n, d_t))

        for (dt, dx, dz) in [(0, 0, 0), (1, 1, 1), (5, 5, 5), (10, 10, 10), (3, 3, 10), (10, 10, 3)]:
            np2sls = NonparametricTwoStageLeastSquares(HermiteFeatures(
                dt), HermiteFeatures(dx), HermiteFeatures(dz), HermiteFeatures(dt, shift=1))
            np2sls.fit(y, p, X=x, W=w, Z=z)
            effect = np2sls.effect(x_fresh, np.zeros(shape(p_fresh)), p_fresh)
            losses.append(np.mean(np.square(p_fresh * x_fresh - effect)))
            marg_effs.append(np2sls.marginal_effect(np.array([[0.3], [0.5], [0.7]]), np.array([[0.4], [0.6], [0.2]])))
        print("losses: {}".format(losses))
        print("marg_effs: {}".format(marg_effs))
