import pytest
import pandas as pd
from primrose.configuration.configuration import Configuration
from primrose.models.sklearn_classifier_model import SklearnClassifierModel
from primrose.data_object import DataObject, DataObjectResponseType
from primrose.readers.csv_reader import CsvReader
from primrose.pipelines.encode_train_test_split import EncodeTrainTestSplit


@pytest.fixture()
def config():
    return Configuration("test/hello_world_tennis.json")


@pytest.fixture()
def data_obj(config):
    df = pd.read_csv("test/tennis.csv")
    data_object = DataObject(config)
    reader = CsvReader(config, 'read_data')
    data_object.add(reader, df)
    encoder = EncodeTrainTestSplit(config, 'encode_and_split')
    data_object, terminate = encoder.run(data_object)
    return data_object


@pytest.fixture
def model(config):
    return SklearnClassifierModel(config, "decision_tree_model")


def test_necessary_config():
    assert SklearnClassifierModel.necessary_config({}) == set(["model_parameters", "mode", "sklearn_classifier_name",
                                                             "grid_search_scoring", "cv_folds"])


def test_get_data(model, data_obj):
    x_train, y_train, x_test, y_test = model._get_data(data_obj)
    assert x_train.shape[0] == y_train.shape[0]
    assert x_test.shape[0] == y_test.shape[0]  # we removed the ID column by not specifying it in X


def test_train_model(model, data_obj):
    model.train_model(data_obj)
    assert model.model is not None
    x_train, y_train, x_test, y_test = model._get_data(data_obj)
    predictions = model.model.predict(x_test)
    assert list(predictions) == [1, 0, 0, 1, 1]


def test_eval(model, data_obj):
    model.train_model(data_obj)
    model.eval_model(data_obj)
    print(model.scores)
    assert round(model.scores['recall'], 2) == 0.67
    assert round(model.scores['precision'], 2) == 0.67
    assert model.scores['predicted_class_fraction'] == 0.6
    assert model.scores['positive_class_fraction'] == 0.6


def test_predict(model, data_obj):
    model.train_model(data_obj)
    data_object = model.predict(data_obj)
    predicted = data_object.get('decision_tree_model', rtype=DataObjectResponseType.KEY_VALUE.value)

    assert predicted['predictions'].shape[0] == 5
    assert list(predicted['predictions'].predictions) == ['yes', 'no', 'no', 'yes', 'yes']
