# -*- coding: utf-8 -*-
"""
Created on Fri Sep 14 11:49:10 2018

@author: Lionel Massoulard
"""

import pytest

import numpy as np
import pandas as pd

from sklearn.base import is_regressor, is_classifier
from sklearn.exceptions import NotFittedError
from sklearn.model_selection import KFold

from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.dummy import DummyRegressor


from aikit.models.stacking import OutSamplerTransformer, StackerClassifier, StackerRegressor


def test_OutSamplerTransformer_classifier():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = 1 * (np.random.randn(100) > 0)

    model = OutSamplerTransformer(RandomForestClassifier(n_estimators=10, random_state=123))
    model.fit(X, y)

    p1 = model.model.predict_proba(X)
    p2 = model.transform(X)

    assert not is_classifier(model)
    assert not is_regressor(model)

    assert np.abs(p1[:, 1] - p2[:, 0]).max() <= 10 ** (-10)
    assert p2.shape == (100, 1)

    assert model.get_feature_names() == ["RandomForestClassifier__1"]

    y = np.array(["a", "b", "c"])[np.random.randint(0, 3, 100)]

    model = OutSamplerTransformer(RandomForestClassifier(n_estimators=10, random_state=123))
    model.fit(X, y)

    p1 = model.model.predict_proba(X)
    p2 = model.transform(X)

    assert p1.shape == (100, 3)
    assert p2.shape == (100, 3)

    assert np.abs(p1 - p2).max() <= 10 ** (-10)

    assert model.get_feature_names() == [
        "RandomForestClassifier__a",
        "RandomForestClassifier__b",
        "RandomForestClassifier__c",
    ]


def test_OutSampleTransformer_classifier_unbalanced():
    np.random.seed(123)
    X = np.random.randn(100, 2)
    y = np.array(["AA"] * 33 + ["BB"] * 33 + ["CC"] * 33 + ["DD"])

    model = OutSamplerTransformer(RandomForestClassifier(n_estimators=10, random_state=123))

    p3 = model.fit_transform(X, y)

    assert (p3.max(axis=1) > 0).all()


def test_OutSamplerTransformer_classifier_fit_transform():

    X = np.random.randn(100, 10)
    y = 1 * (np.random.randn(100) > 0)

    cv = KFold(n_splits=10, shuffle=True, random_state=123)

    model = OutSamplerTransformer(LogisticRegression(C=1,random_state=123), cv=cv)
    model.fit(X, y)
    y1 = model.transform(X)

    model = OutSamplerTransformer(LogisticRegression(C=1,random_state=123), cv=cv)
    y2 = model.fit_transform(X, y)

    assert np.abs(y1 - y2).flatten().max() >= 0.01  # vector should be different


def test_OutSamplerTransformer_regressor():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = np.random.randn(100)

    model = OutSamplerTransformer(RandomForestRegressor(n_estimators=10,random_state=123), cv=10)
    model.fit(X, y)

    y1 = model.model.predict(X)
    y2 = model.transform(X)

    assert not is_classifier(model)
    assert not is_regressor(model)

    assert np.abs(y1 - y2[:, 0]).max() <= 10 ** (-10)
    assert y2.shape == (100, 1)

    assert model.get_feature_names() == ["RandomForestRegressor__target"]


def test_OutSamplerTransformer_regressor_fit_transform():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = np.random.randn(100)

    cv = KFold(n_splits=10, shuffle=True, random_state=123)

    model = OutSamplerTransformer(DummyRegressor(), cv=cv)
    model.fit(X, y)
    y1 = model.transform(X)

    model = OutSamplerTransformer(DummyRegressor(), cv=cv)
    y2 = model.fit_transform(X, y)

    assert np.abs(y1 - y2).flatten().max() >= 0.01  # vector should be different


def test_approx_cross_validation_OutSamplerTransformer_regressor():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = np.random.randn(100)

    model = OutSamplerTransformer(RandomForestRegressor(random_state=123), cv=10)

    cv_res, yhat = model.approx_cross_validation(X, y, cv=10, method="transform", no_scoring=True)

    assert cv_res is None
    assert yhat.ndim == 2
    assert yhat.shape == (y.shape[0], 1)

    with pytest.raises(NotFittedError):
        model.transform(X)

    cv = KFold(n_splits=10, shuffle=True, random_state=123)

    model = OutSamplerTransformer(DummyRegressor(), cv=cv)
    yhat1 = model.fit_transform(X, y)

    cv_res, yhat2 = model.approx_cross_validation(X, y, cv=cv, method="transform", no_scoring=True, return_predict=True)
    # Approx cross val and fit transform should return the same thing here
    assert np.abs((yhat1 - yhat2).flatten()).max() <= 10 ** (-5)

    yhat3 = np.zeros((y.shape[0], 1))

    for train, test in cv.split(X, y):
        model = DummyRegressor()
        model.fit(X[train, :], y[train])

        yhat3[test, 0] = model.predict(X[test, :])

    assert np.abs((yhat1 - yhat3).flatten()).max() <= 10 ** (-5)
    assert np.abs((yhat1 - yhat2).flatten()).max() <= 10 ** (-5)


def test_approx_cross_validation_OutSamplerTransformer_classifier():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = 1 * (np.random.randn(100) > 0)

    model = OutSamplerTransformer(RandomForestClassifier(random_state=123), cv=10)

    cv_res, yhat = model.approx_cross_validation(X, y, cv=10, method="transform", no_scoring=True)

    assert cv_res is None
    assert yhat.ndim == 2
    assert yhat.shape == (y.shape[0], 1)

    with pytest.raises(NotFittedError):
        model.transform(X)

    with pytest.raises(NotFittedError):
        model.model.predict(X)

    cv = KFold(n_splits=10, shuffle=True, random_state=123)
    model = OutSamplerTransformer(LogisticRegression(C=1,random_state=123), cv=cv)
    yhat1 = model.fit_transform(X, y)

    model = OutSamplerTransformer(LogisticRegression(C=1,random_state=123), cv=cv)
    cv_res, yhat2 = model.approx_cross_validation(X, y, cv=cv, method="transform", no_scoring=True, return_predict=True)

    # Approx cross val and fit transform should return the same thing here
    assert np.abs((yhat1 - yhat2).flatten()).max() <= 10 ** (-5)

    yhat3 = np.zeros((y.shape[0], 1))

    for train, test in cv.split(X, y):
        model = LogisticRegression()
        model.fit(X[train, :], y[train])

        yhat3[test, 0] = model.predict_proba(X[test, :])[:, 1]

    assert np.abs((yhat1 - yhat3).flatten()).max() <= 10 ** (-5)
    assert np.abs((yhat1 - yhat2).flatten()).max() <= 10 ** (-5)


def test_StackerRegressor():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = np.random.randn(100)

    stacker = StackerRegressor(models=[RandomForestRegressor(n_estimators=10,random_state=123), Ridge(random_state=123)], cv=10, blender=Ridge(random_state=123))

    stacker.fit(X, y)

    yhat = stacker.predict(X)

    assert yhat.ndim == 1
    assert yhat.shape[0] == X.shape[0]

    assert is_regressor(stacker)
    assert not is_classifier(stacker)

    with pytest.raises(AttributeError):
        stacker.predict_proba(X)

    with pytest.raises(AttributeError):
        stacker.classes_


def test_StackerClassifier():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = 1 * (np.random.randn(100) > 0)

    stacker = StackerClassifier(
        models=[RandomForestClassifier(random_state=123), LogisticRegression(C=1,random_state=123)], cv=10, blender=LogisticRegression(C=1,random_state=123)
    )

    stacker.fit(X, y)

    yhat = stacker.predict(X)

    assert yhat.ndim == 1
    assert yhat.shape[0] == X.shape[0]

    assert list(set(yhat)) == [0, 1]
    assert list(stacker.classes_) == [0, 1]

    yhat_proba = stacker.predict_proba(X)

    assert yhat_proba.shape == (y.shape[0], 2)

    assert not is_regressor(stacker)
    assert is_classifier(stacker)


def test_approx_cross_validation_StackerRegressor():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = np.random.randn(100)

    stacker = StackerRegressor(models=[RandomForestRegressor(n_estimators=10,random_state=123), Ridge(random_state=123)], cv=10, blender=Ridge(random_state=123))

    cv_res, yhat = stacker.approx_cross_validation(
        X, y, cv=10, method="predict", scoring=["neg_mean_squared_error"], return_predict=True, verbose=False
    )

    assert cv_res is not None
    assert isinstance(cv_res, pd.DataFrame)
    assert cv_res.shape[0] == 10
    assert "test_neg_mean_squared_error" in cv_res
    assert "train_neg_mean_squared_error" in cv_res

    assert yhat.ndim == 1
    assert yhat.shape[0] == y.shape[0]

    with pytest.raises(NotFittedError):
        stacker.predict(X)

    for m in stacker.models:
        with pytest.raises(NotFittedError):
            m.predict(X)


def test_approx_cross_validation_StackerClassifier():

    np.random.seed(123)
    X = np.random.randn(100, 10)
    y = 1 * (np.random.randn(100) > 0)

    stacker = StackerClassifier(
        models=[RandomForestClassifier(n_estimators=10,random_state=123), LogisticRegression(C=1,random_state=123)], cv=10, blender=LogisticRegression(C=1,random_state=123)
    )

    cv_res, yhat = stacker.approx_cross_validation(
        X, y, cv=10, method="predict_proba", scoring=["accuracy"], return_predict=True, verbose=False
    )

    assert cv_res is not None
    assert isinstance(cv_res, pd.DataFrame)
    assert cv_res.shape[0] == 10
    assert "test_accuracy" in cv_res
    assert "train_accuracy" in cv_res

    assert yhat.ndim == 2
    assert yhat.shape == (y.shape[0], 2)

    with pytest.raises(NotFittedError):
        stacker.predict(X)

    for m in stacker.models:
        with pytest.raises(NotFittedError):
            m.predict(X)


def _verif_all():
    test_OutSamplerTransformer_classifier()
    test_OutSamplerTransformer_regressor()

    test_OutSamplerTransformer_classifier_fit_transform()
    test_OutSamplerTransformer_regressor_fit_transform()

    test_approx_cross_validation_OutSamplerTransformer_regressor()
    test_approx_cross_validation_OutSamplerTransformer_classifier()

    test_StackerClassifier()
    test_StackerRegressor()

    test_approx_cross_validation_StackerClassifier()
    test_approx_cross_validation_StackerRegressor()

[(28, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (29, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (67, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (68, 'numpy.random.randn', 'np.random.randn', (['(100)', '(2)'], {}), True, 'import numpy as np\n'), (69, 'numpy.array', 'np.array', (["(['AA'] * 33 + ['BB'] * 33 + ['CC'] * 33 + ['DD'])"], {}), True, 'import numpy as np\n'), (80, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (83, 'sklearn.model_selection.KFold', 'KFold', ([], {'n_splits': '(10)', 'shuffle': '(True)', 'random_state': '(123)'}), False, 'from sklearn.model_selection import KFold\n'), (97, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (98, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (99, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (118, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (119, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (120, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (122, 'sklearn.model_selection.KFold', 'KFold', ([], {'n_splits': '(10)', 'shuffle': '(True)', 'random_state': '(123)'}), False, 'from sklearn.model_selection import KFold\n'), (136, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (137, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (138, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (151, 'sklearn.model_selection.KFold', 'KFold', ([], {'n_splits': '(10)', 'shuffle': '(True)', 'random_state': '(123)'}), False, 'from sklearn.model_selection import KFold\n'), (160, 'numpy.zeros', 'np.zeros', (['(y.shape[0], 1)'], {}), True, 'import numpy as np\n'), (174, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (175, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (192, 'sklearn.model_selection.KFold', 'KFold', ([], {'n_splits': '(10)', 'shuffle': '(True)', 'random_state': '(123)'}), False, 'from sklearn.model_selection import KFold\n'), (202, 'numpy.zeros', 'np.zeros', (['(y.shape[0], 1)'], {}), True, 'import numpy as np\n'), (216, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (217, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (218, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (229, 'sklearn.base.is_regressor', 'is_regressor', (['stacker'], {}), False, 'from sklearn.base import is_regressor, is_classifier\n'), (241, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (242, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (264, 'sklearn.base.is_classifier', 'is_classifier', (['stacker'], {}), False, 'from sklearn.base import is_regressor, is_classifier\n'), (269, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (270, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (271, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (298, 'numpy.random.seed', 'np.random.seed', (['(123)'], {}), True, 'import numpy as np\n'), (299, 'numpy.random.randn', 'np.random.randn', (['(100)', '(10)'], {}), True, 'import numpy as np\n'), (32, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', ([], {'n_estimators': '(10)', 'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestClassifier\n'), (38, 'sklearn.base.is_classifier', 'is_classifier', (['model'], {}), False, 'from sklearn.base import is_regressor, is_classifier\n'), (39, 'sklearn.base.is_regressor', 'is_regressor', (['model'], {}), False, 'from sklearn.base import is_regressor, is_classifier\n'), (46, 'numpy.array', 'np.array', (["['a', 'b', 'c']"], {}), True, 'import numpy as np\n'), (48, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', ([], {'n_estimators': '(10)', 'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestClassifier\n'), (71, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', ([], {'n_estimators': '(10)', 'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestClassifier\n'), (85, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {'C': '(1)', 'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (89, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {'C': '(1)', 'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (101, 'sklearn.ensemble.RandomForestRegressor', 'RandomForestRegressor', ([], {'n_estimators': '(10)', 'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestRegressor\n'), (107, 'sklearn.base.is_classifier', 'is_classifier', (['model'], {}), False, 'from sklearn.base import is_regressor, is_classifier\n'), (108, 'sklearn.base.is_regressor', 'is_regressor', (['model'], {}), False, 'from sklearn.base import is_regressor, is_classifier\n'), (124, 'sklearn.dummy.DummyRegressor', 'DummyRegressor', ([], {}), False, 'from sklearn.dummy import DummyRegressor\n'), (128, 'sklearn.dummy.DummyRegressor', 'DummyRegressor', ([], {}), False, 'from sklearn.dummy import DummyRegressor\n'), (140, 'sklearn.ensemble.RandomForestRegressor', 'RandomForestRegressor', ([], {'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestRegressor\n'), (148, 'pytest.raises', 'pytest.raises', (['NotFittedError'], {}), False, 'import pytest\n'), (153, 'sklearn.dummy.DummyRegressor', 'DummyRegressor', ([], {}), False, 'from sklearn.dummy import DummyRegressor\n'), (163, 'sklearn.dummy.DummyRegressor', 'DummyRegressor', ([], {}), False, 'from sklearn.dummy import DummyRegressor\n'), (178, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', ([], {'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestClassifier\n'), (186, 'pytest.raises', 'pytest.raises', (['NotFittedError'], {}), False, 'import pytest\n'), (189, 'pytest.raises', 'pytest.raises', (['NotFittedError'], {}), False, 'import pytest\n'), (193, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {'C': '(1)', 'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (196, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {'C': '(1)', 'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (205, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (230, 'sklearn.base.is_classifier', 'is_classifier', (['stacker'], {}), False, 'from sklearn.base import is_regressor, is_classifier\n'), (232, 'pytest.raises', 'pytest.raises', (['AttributeError'], {}), False, 'import pytest\n'), (235, 'pytest.raises', 'pytest.raises', (['AttributeError'], {}), False, 'import pytest\n'), (263, 'sklearn.base.is_regressor', 'is_regressor', (['stacker'], {}), False, 'from sklearn.base import is_regressor, is_classifier\n'), (288, 'pytest.raises', 'pytest.raises', (['NotFittedError'], {}), False, 'import pytest\n'), (319, 'pytest.raises', 'pytest.raises', (['NotFittedError'], {}), False, 'import pytest\n'), (30, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (46, 'numpy.random.randint', 'np.random.randint', (['(0)', '(3)', '(100)'], {}), True, 'import numpy as np\n'), (81, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (176, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (220, 'sklearn.linear_model.Ridge', 'Ridge', ([], {'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (243, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (246, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {'C': '(1)', 'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (273, 'sklearn.linear_model.Ridge', 'Ridge', ([], {'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (292, 'pytest.raises', 'pytest.raises', (['NotFittedError'], {}), False, 'import pytest\n'), (300, 'numpy.random.randn', 'np.random.randn', (['(100)'], {}), True, 'import numpy as np\n'), (303, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {'C': '(1)', 'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (323, 'pytest.raises', 'pytest.raises', (['NotFittedError'], {}), False, 'import pytest\n'), (41, 'numpy.abs', 'np.abs', (['(p1[:, (1)] - p2[:, (0)])'], {}), True, 'import numpy as np\n'), (57, 'numpy.abs', 'np.abs', (['(p1 - p2)'], {}), True, 'import numpy as np\n'), (110, 'numpy.abs', 'np.abs', (['(y1 - y2[:, (0)])'], {}), True, 'import numpy as np\n'), (220, 'sklearn.ensemble.RandomForestRegressor', 'RandomForestRegressor', ([], {'n_estimators': '(10)', 'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestRegressor\n'), (220, 'sklearn.linear_model.Ridge', 'Ridge', ([], {'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (246, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', ([], {'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestClassifier\n'), (246, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {'C': '(1)', 'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (273, 'sklearn.ensemble.RandomForestRegressor', 'RandomForestRegressor', ([], {'n_estimators': '(10)', 'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestRegressor\n'), (273, 'sklearn.linear_model.Ridge', 'Ridge', ([], {'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (303, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', ([], {'n_estimators': '(10)', 'random_state': '(123)'}), False, 'from sklearn.ensemble import RandomForestClassifier\n'), (303, 'sklearn.linear_model.LogisticRegression', 'LogisticRegression', ([], {'C': '(1)', 'random_state': '(123)'}), False, 'from sklearn.linear_model import LogisticRegression, Ridge\n'), (92, 'numpy.abs', 'np.abs', (['(y1 - y2)'], {}), True, 'import numpy as np\n'), (131, 'numpy.abs', 'np.abs', (['(y1 - y2)'], {}), True, 'import numpy as np\n')]