from unittest.mock import ANY, MagicMock, patch
from pandas import DataFrame
from pytest import fixture, mark

from laborious.activities.gates import Gates
from sientia_do.notifications.models import NotificationLevel


@fixture
def gates():
    return Gates(
        logger=MagicMock(),
        notification_handler=MagicMock()
    )


@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_specific_variables_null_values_with_stop_policy_only(
    input_filter_functions_mock,
    gates
):
    specific_variables_null_values_mock = MagicMock(return_value=True)
    empty_data_mock = MagicMock(return_value=False)

    def functions_side_effect(x):
        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
            return specific_variables_null_values_mock
        if x == 'path_confidence':
            return {
                'stop': -1,
                'continue': 2,
                'repeat': -1
            }
        return empty_data_mock

    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect

    input_data = {
        'filters': {
            'SPECIFIC_VARIABLES_NULL_VALUES': {
                'POLICY': 'stop',
                'VARIABLES': ['variable2']
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat']
    }

    result = await gates.input_gate(input_data)
    assert result == ('stop', -1)

    input_args = specific_variables_null_values_mock.call_args
    assert input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert input_args[0][1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']

    empty_data_mock.assert_not_called()


@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_specific_variables_null_values_with_continue_policy_only(
    input_filter_functions_mock,
    gates
):
    specific_variables_null_values_mock = MagicMock(return_value=True)
    empty_data_mock = MagicMock(return_value=False)

    def functions_side_effect(x):
        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
            return specific_variables_null_values_mock
        if x == 'path_confidence':
            return {
                'stop': -1,
                'continue': 2,
                'repeat': -1
            }
        return empty_data_mock

    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect

    input_data = {
        'filters': {
            'SPECIFIC_VARIABLES_NULL_VALUES': {
                'POLICY': 'continue',
                'VARIABLES': ['variable2']
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat']
    }

    result = await gates.input_gate(input_data)
    assert result == ('continue', 2)

    input_args = specific_variables_null_values_mock.call_args
    assert input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert input_args[0][1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']

    empty_data_mock.assert_not_called()


@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_specific_variables_null_values_no_filtered(
    input_filter_functions_mock,
    gates
):
    specific_variables_null_values_mock = MagicMock(return_value=False)
    empty_data_mock = MagicMock(return_value=False)

    def functions_side_effect(x):
        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
            return specific_variables_null_values_mock
        if x == 'path_confidence':
            return {
                'stop': -1,
                'continue': 2,
                'repeat': -1
            }
        return empty_data_mock

    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect

    input_data = {
        'filters': {
            'SPECIFIC_VARIABLES_NULL_VALUES': {
                'POLICY': 'stop',
                'VARIABLES': ['variable2']
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat']
    }

    result = await gates.input_gate(input_data)
    assert result == (None, 0)

    specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args

    assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert specific_variables_null_values_input_args[0][
        1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']

    empty_data_mock.assert_not_called()


@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_one_stop_policy(
    input_filter_functions_mock,
    gates
):
    specific_variables_null_values_mock = MagicMock(return_value=True)
    empty_data_mock = MagicMock(return_value=True)

    def functions_side_effect(x):
        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
            return specific_variables_null_values_mock
        if x == 'path_confidence':
            return {
                'stop': -1,
                'continue': 2,
                'repeat': -1
            }
        return empty_data_mock

    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect

    input_data = {
        'filters': {
            'SPECIFIC_VARIABLES_NULL_VALUES': {
                'POLICY': 'stop',
                'VARIABLES': ['variable2']
            },
            'EMPTY_DATA': {
                'POLICY': 'continue',
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat']
    }

    result = await gates.input_gate(input_data)
    assert result == ('stop', -1)

    specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
    assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert specific_variables_null_values_input_args[0][
        1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']

    empty_data_input_args = empty_data_mock.call_args
    assert empty_data_input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']


@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_one_continue_policy(
    input_filter_functions_mock,
    gates
):
    specific_variables_null_values_mock = MagicMock(return_value=False)
    empty_data_mock = MagicMock(return_value=True)

    def functions_side_effect(x):
        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
            return specific_variables_null_values_mock
        if x == 'path_confidence':
            return {
                'stop': -1,
                'continue': 2,
                'repeat': -1
            }
        return empty_data_mock

    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect

    input_data = {
        'filters': {
            'SPECIFIC_VARIABLES_NULL_VALUES': {
                'POLICY': 'stop',
                'VARIABLES': ['variable2']
            },
            'EMPTY_DATA': {
                'POLICY': 'continue',
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat']
    }

    result = await gates.input_gate(input_data)
    assert result == ('continue', 2)

    specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
    assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert specific_variables_null_values_input_args[0][
        1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']

    empty_data_input_args = empty_data_mock.call_args
    assert empty_data_input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']


@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_no_filtered(
    input_filter_functions_mock,
    gates
):
    specific_variables_null_values_mock = MagicMock(return_value=False)
    empty_data_mock = MagicMock(return_value=False)

    def functions_side_effect(x):
        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
            return specific_variables_null_values_mock
        if x == 'path_confidence':
            return {
                'stop': -1,
                'continue': 2,
                'repeat': -1
            }
        return empty_data_mock

    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect

    input_data = {
        'filters': {
            'SPECIFIC_VARIABLES_NULL_VALUES': {
                'POLICY': 'stop',
                'VARIABLES': ['variable2']
            },
            'EMPTY_DATA': {
                'POLICY': 'continue',
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat']
    }

    result = await gates.input_gate(input_data)
    assert result == (None, 0)

    specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
    assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))

    empty_data_input_args = empty_data_mock.call_args
    assert empty_data_input_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']


@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_error(
    input_filter_functions_mock,
    gates
):
    input_filter_functions_mock.__getitem__.side_effect = KeyError('test')

    input_data = {
        'filters': {
            'SPECIFIC_VARIABLES_NULL_VALUES': {
                'POLICY': 'stop',
                'VARIABLES': ['variable2']
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat'],
    }

    result = await gates.input_gate(input_data)
    assert result == (None, 0)

    gates.notification_handler.build_and_send_notification.assert_called_once_with(
        notification_id='INTPUT_GATE_ERROR__SPECIFIC_VARIABLES_NULL_VALUES',
        message="Error in filter SPECIFIC_VARIABLES_NULL_VALUES:{'POLICY': 'stop', 'VARIABLES': ['variable2']}: \n 'test'",
        block='input_gate',
        level=NotificationLevel.ERROR,
        attachment_content=ANY
    )


transform_filter_path_confidence = {
    'stop': -1,
    'continue': 255,
    'repeat': -1
}


@mark.asyncio
@patch('laborious.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_no_filtered(
    mlflow_response_filter_functions_mock,
    gates
):
    api_error_filter_mock = MagicMock(return_value=False)

    def transform_filter_functions_side_effect(x: str):
        if x == 'API_ERROR':
            return api_error_filter_mock
        if x == 'path_confidence':
            return transform_filter_path_confidence

    mlflow_response_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect

    input_data = {
        'filters': {
            'API_ERROR': {
                'POLICY': 'stop',
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat'],
        'type': 'predict'
    }

    result = await gates.mlflow_response_gate(input_data)
    assert result == (None, 0)

    api_error_filter_mock.assert_called_once_with(
        input_data['data'],
        input_data['filters']['API_ERROR']
    )

    gates.notification_handler.build_and_send_notification.assert_not_called()


@mark.asyncio
@patch('laborious.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_filtered(
    mlflow_response_filter_functions_mock,
    gates
):
    api_error_filter_mock = MagicMock(return_value=True)

    def transform_filter_functions_side_effect(x: str):
        if x == 'API_ERROR':
            return api_error_filter_mock
        if x == 'path_confidence':
            return transform_filter_path_confidence

    mlflow_response_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect

    input_data = {
        'filters': {
            'API_ERROR': {
                'POLICY': 'continue',
            }
        },
        'data': {
            'success': False,
            'content': {
                'message': 'Error',
                'traceback': 'Error'
            }
        },
        'path_priority': ['stop', 'continue', 'repeat'],
        'type': 'predict'
    }

    result = await gates.mlflow_response_gate(input_data)
    assert result == ('continue', 255)

    api_error_filter_mock.assert_called_once_with(
        input_data['data'],
        input_data['filters']['API_ERROR']
    )

    gates.notification_handler.build_and_send_notification.assert_called_once_with(
        notification_id='PREDICT_GATE_RESPONSE_FILTER__API_ERROR',
        message=input_data['data']['content']['message'],
        block='mlflow_gate',
        level=NotificationLevel.WARNING,
        attachment_content=input_data['data']['content']['traceback']
    )


@mark.asyncio
@patch('laborious.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_no_filtered(
    mlflow_content_filter_functions_mock,
    gates
):
    nan_values_filter_mock = MagicMock(return_value=False)

    def transform_filter_functions_side_effect(x: str):
        if x == 'NAN_VALUES':
            return nan_values_filter_mock
        if x == 'path_confidence':
            return transform_filter_path_confidence

    mlflow_content_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect

    input_data = {
        'filters': {
            'NAN_VALUES': {
                'POLICY': 'repeat',
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat'],
        'type': 'predict'
    }

    result = await gates.mlflow_content_gate(input_data)
    assert result == (None, 0)

    nan_values_filter_mock_args = nan_values_filter_mock.call_args
    assert nan_values_filter_mock_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert nan_values_filter_mock_args[0][1] == input_data['filters']['NAN_VALUES']

    gates.notification_handler.build_and_send_notification.assert_not_called()


@mark.asyncio
@patch('laborious.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_filtered(
    mlflow_content_filter_functions_mock,
    gates
):
    nan_values_filter_mock = MagicMock(return_value=True)

    def transform_filter_functions_side_effect(x: str):
        if x == 'NAN_VALUES':
            return nan_values_filter_mock
        if x == 'path_confidence':
            return transform_filter_path_confidence

    mlflow_content_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect

    input_data = {
        'filters': {
            'NAN_VALUES': {
                'POLICY': 'repeat',
            }
        },
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'path_priority': ['stop', 'continue', 'repeat'],
        'type': 'predict'
    }

    result = await gates.mlflow_content_gate(input_data)
    assert result == ('repeat', -1)

    nan_values_filter_mock_args = nan_values_filter_mock.call_args
    assert nan_values_filter_mock_args[0][0].equals(DataFrame(
        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
    assert nan_values_filter_mock_args[0][1] == input_data['filters']['NAN_VALUES']

    gates.notification_handler.build_and_send_notification.assert_called_once_with(
        notification_id='PREDICT_GATE_CONTENT_FILTER__NAN_VALUES',
        message="Data not passed the content filter NAN_VALUES:{'POLICY': 'repeat'}",
        block='mlflow_gate',
        level=NotificationLevel.WARNING,
        attachment_content=DataFrame(input_data['data']).to_string()
    )


@mark.asyncio
async def test_format_prediction(
    gates
):
    input_data = {
        'data': {
            'variable': ['variable1', 'variable2'],
            'value': [1, 2]
        },
        'timestamp': '2021-01-01',
        'model_id': 'model_id',
        'prediction_confidence': 0.95
    }

    expected_output = DataFrame(input_data['data'])
    expected_output['timestamp'] = input_data['timestamp']
    expected_output['model_id'] = input_data['model_id']
    expected_output['prediction_confidence'] = input_data['prediction_confidence']
    expected_output['prediction_status'] = 'Good'
    expected_output['comment'] = ''

    result = await gates.format_prediction(input_data)
    assert result == expected_output.to_dict()


@mark.asyncio
async def test_format_default_prediction(
    gates
):
    input_data = {
        'timestamp': '2021-01-01',
        'model_id': 'model_id',
        'prediction_confidence': 0.95,
        'comment': 'Comment'
    }

    expected_output = DataFrame({
        'prediction': [0],
        'response_time': [0],
        'timestamp': [input_data['timestamp']],
        'model_id': [input_data['model_id']],
        'prediction_confidence': [input_data['prediction_confidence']],
        'prediction_status': ['Bad'],
        'comment': [input_data['comment']]
    })

    result = await gates.format_default_prediction(input_data)
    assert result == expected_output.to_dict()