import numpy as np

import warnings

import pytest

from triage.component.catwalk.estimators.transformers import CutOff
from triage.component.catwalk.estimators.classifiers import ScaledLogisticRegression

from sklearn import linear_model

from sklearn import datasets
from sklearn import preprocessing
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split

@pytest.fixture
def data():
    dataset = datasets.load_breast_cancer()
    X = dataset.data
    y = dataset.target

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.3, random_state=12345)

    return {'X_train':X_train, 'X_test':X_test, 'y_train':y_train, 'y_test':y_test}

def test_cutoff_warning():
    X_data = [1, 2, 0.5, 0.7, 100, -1, -23, 0]

    cutoff = CutOff()

    with pytest.warns(UserWarning):
       cutoff.fit_transform(X_data)


def test_cutoff_transformer():
    cutoff = CutOff()

    X_data = [1, 2, 0.5, 0.7, 100, -1, -23, 0]

    assert np.all(cutoff.fit_transform(X_data) == [1, 1, 0.5, 0.7, 1, 0, 0, 0])

def test_cutoff_inside_a_pipeline(data):
    minmax_scaler = preprocessing.MinMaxScaler()
    dsapp_cutoff = CutOff()

    pipeline =Pipeline([
        ('minmax_scaler',minmax_scaler),
        ('dsapp_cutoff', dsapp_cutoff)
    ])

    pipeline.fit(data['X_train'], data['y_train'])

    X_fake_new_data = data['X_test'][-1,:].reshape(1,-1) + 0.5

    mms = preprocessing.MinMaxScaler().fit(data['X_train'])

    assert np.all(( mms.transform(X_fake_new_data) > 1  ) == (pipeline.transform(X_fake_new_data) == 1))

def test_dsapp_lr(data):
    dsapp_lr = ScaledLogisticRegression()
    dsapp_lr.fit(data['X_train'], data['y_train'])

    minmax_scaler = preprocessing.MinMaxScaler()
    dsapp_cutoff = CutOff()
    lr = linear_model.LogisticRegression()

    pipeline =Pipeline([
        ('minmax_scaler',minmax_scaler),
        ('dsapp_cutoff', dsapp_cutoff),
        ('lr', lr)
    ])

    pipeline.fit(data['X_train'], data['y_train'])

    assert np.all(dsapp_lr.predict(data['X_test']) == pipeline.predict(data['X_test']))
