#!/usr/bin/env python
# coding=utf-8

try:
    import mock

except ImportError:
    import unittest.mock as mock

import pandas as pd
from marvin_titanic_engine.training import MetricsEvaluator


@mock.patch('marvin_titanic_engine.training.metrics_evaluator.metrics.confusion_matrix')
@mock.patch('marvin_titanic_engine.training.metrics_evaluator.metrics.classification_report')
def test_execute(report_mocked, matrix_mocked, mocked_params):

    test_dataset = {
        "X_test": pd.DataFrame({'Sex': [1, 2, 3], 'B': [4, 5, 6], 'C': [7, 8, 9]}),
        "y_test": pd.DataFrame({'Sex': [1, 2, 3], 'B': [4, 5, 6], 'C': [7, 8, 9]}),
        "X_train": pd.DataFrame({'Sex': [1, 2, 3], 'B': [4, 5, 6], 'C': [7, 8, 9]}),
        "y_train": pd.DataFrame({'Sex': [1, 2, 3], 'B': [4, 5, 6], 'C': [7, 8, 9]}),
        "sss": mock.MagicMock()
    }

    mocked_params = {
        "pred_cols": ["Sex", "B"],
        "dep_var": "C"
    }

    model_mocked = {
        # "model_type": "test_type",
        "test": mock.MagicMock(),
        "rf": mock.MagicMock()
    }

    ac = MetricsEvaluator(model=model_mocked, dataset=test_dataset)
    ac.execute(params=mocked_params)

    report_mocked.assert_called()
    matrix_mocked.assert_called()
    model_mocked["test"].predict.assert_called_once()