import csv
import io
import json
import six
import tempfile

import pandas as pd
import pytest
import responses
import trafaret

from datarobot import TrainingPredictions
from datarobot.models.training_predictions import TrainingPredictionsRow


def _write_csv(lines):

    if six.PY2:

        buf = io.BytesIO()
        wr = csv.writer(buf)
        wr.writerows(lines)
        return buf.getvalue()

    buf = io.StringIO()
    wr = csv.writer(buf)
    wr.writerows(lines)
    return buf.getvalue().encode('utf-8')


@pytest.fixture
def prediction_id():
    return '59d92fe8962d7464ca8c6bd6'


@pytest.fixture
def predictions_with_shap_exp_csv():
    return """row_id,partition_id,prediction,class_1.0,class_0.0,explained_class,Explanation_1_feature_name,Explanation_1_feature_value,Explanation_1_strength,shap_remaining_total,shap_base_value
        0,TrainTestOverlap,0.0,0.311178247734139,0.6888217522658611,1.0,number_diagnoses,9,-0.07643818895488698,0.006562548084582246,0.39640626311300003
        1,TrainTestOverlap,0.0,0.2829861111111111,0.7170138888888888,1.0,number_diagnoses,6,0.06648347981438553,-0.013534302788801085,0.39640626311300003
        2,TrainTestOverlap,0.0,0.42699724517906334,0.5730027548209367,1.0,number_diagnoses,9,0.08556893050818416,-0.032619753482599656,0.39640626311300003
        3,TrainTestOverlap,0.0,0.32653061224489793,0.6734693877551021,1.0,number_diagnoses,3,0.05807900167179779,-0.0051298246462132885,0.39640626311300003
        4,TrainTestOverlap,0.0,0.06,0.94,1.0,number_inpatient,0,-0.056383896138088974,-0.0570362457984236,0.39640626311300003
        5,TrainTestOverlap,0.0,0.3774597495527728,0.6225402504472273,1.0,number_diagnoses,8,0.04847090471521205,-0.03250602234256627,0.39640626311300003
        6,TrainTestOverlap,0.0,0.13768115942028986,0.8623188405797102,1.0,admission_type_id,Elective,0.08996873371035077,0.040252262429918834,0.39640626311300003
        7,TrainTestOverlap,1.0,0.5300353356890459,0.4699646643109541,1.0,admission_type_id,Emergency,0.16263335056706338,-0.03241235442679378,0.39640626311300003
        8,TrainTestOverlap,1.0,0.6408668730650154,0.35913312693498456,1.0,number_inpatient,1,0.34801989552094514,-0.055934082238188985,0.39640626311300003
        9,TrainTestOverlap,0.0,0.2027972027972028,0.7972027972027972,1.0,number_inpatient,0,0.1146691220002886,-0.08407812406711129,0.39640626311300003
    """


@responses.activate
def test_get__no_requests_made(project_id, prediction_id):
    TrainingPredictions.get(project_id, prediction_id)
    assert len(responses.calls) == 0


@responses.activate
def test_list__ok(project_id, model_id, prediction_id):
    training_prediction_stub = [
        {
            'url': 'https://host_name.com/projects/{}/trainingPredictions/{}//'.format(
                project_id, prediction_id
            ),
            'id': prediction_id,
            'modelId': model_id,
            'dataSubset': 'all',
        },
        {
            'url': 'https://host_name.com/projects/{}/trainingPredictions/{}//'.format(
                project_id, prediction_id
            ),
            'id': prediction_id,
            'modelId': model_id,
            'dataSubset': 'all',
            'explanationAlgorithm': 'shap',
            'maxExplanations': 7,
            'shapWarnings': {
                'mismatchRowCount': 1,
                'maxNormalizedMismatch': 0.01,
            }
        }
    ]
    responses.add(
        responses.GET,
        'https://host_name.com/projects/{}/trainingPredictions/'.format(
            project_id,
        ),
        status=200,
        body=json.dumps({'data': training_prediction_stub}),
    )

    training_predictions_list = TrainingPredictions.list(project_id)

    assert len(training_predictions_list) == 2
    training_prediction = training_predictions_list[0]
    assert training_prediction.project_id == project_id
    assert training_prediction.model_id == model_id
    assert training_prediction.prediction_id == prediction_id
    training_prediction_with_shap_exp = training_predictions_list[1]
    assert training_prediction_with_shap_exp.project_id == project_id
    assert training_prediction_with_shap_exp.model_id == model_id
    assert training_prediction_with_shap_exp.explanation_algorithm == 'shap'
    assert training_prediction_with_shap_exp.max_explanations == 7
    assert training_prediction_with_shap_exp.shap_warnings == {
        'mismatch_row_count': 1,
        'max_normalized_mismatch': 0.01,
    }


@responses.activate
def test_list__no_url_in_api_response__error(project_id, model_id, prediction_id):
    training_prediction_stub = {
        'url': '',
        'id': prediction_id,
        'modelId': model_id,
        'trainingPredictionSet': 'all',
    }
    responses.add(
        responses.GET,
        'https://host_name.com/projects/{}/trainingPredictions/'.format(
            project_id,
        ),
        status=200,
        body=json.dumps({'data': [training_prediction_stub]}),
    )

    with pytest.raises(trafaret.DataError):
        TrainingPredictions.list(project_id)


@responses.activate
def test_iterate_rows__empty_predictions__ok(project_id, prediction_id):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    response = {'data': [], 'next': '{}?offset=100&limit=100'.format(url)}
    responses.add(
        responses.GET,
        url,
        body=json.dumps(response),
        status=200,
        content_type='application/json',
    )

    obj = TrainingPredictions.get(project_id, prediction_id)

    assert list(obj.iterate_rows()) == []


@responses.activate
def test_iterate_rows__malfromed_predictions__data_error(project_id, prediction_id):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    items = [{
        'row_id': 'wrong-id',
        'partition_id': 'holdout',
        'prediction': 0.42,
        'prediction_values': [],
    }]
    response = {'data': items, 'next': '{}?offset=100&limit=100'.format(url)}
    responses.add(
        responses.GET,
        url,
        body=json.dumps(response),
        status=200,
        content_type='application/json',
    )

    with pytest.raises(trafaret.DataError):
        obj = TrainingPredictions.get(project_id, prediction_id)
        it = obj.iterate_rows()
        next(it)


@responses.activate
def test_iterate_rows__extra_fields__ignored(project_id, prediction_id):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    items = [{
        'row_id': 1,
        'partition_id': 2,
        'prediction': 0.42,
        'prediction_values': [{
            'label': 'test-label',
            'value': 0.42,
            'extra-prediction': 'you are going to win a lottery for sure',
        }]},
    ]
    response = {'data': items, 'next': '{}?offset=100&limit=100'.format(url), 'prev': None}
    responses.add(
        responses.GET,
        url,
        body=json.dumps(response),
        status=200,
        content_type='application/json',
    )

    obj = TrainingPredictions.get(project_id, prediction_id)
    assert list(obj.iterate_rows()) == [TrainingPredictionsRow(
        row_id=1,
        partition_id=2,
        prediction=0.42,
        prediction_values=[{
            'label': 'test-label',
            'value': 0.42,
        }],
        timestamp=None,
        forecast_point=None,
        forecast_distance=None,
        series_id=None,
        prediction_explanations=None,
        shap_metadata=None,
    )]


@responses.activate
def test_iterate_rows__all_rows_at_once__ok(project_id, prediction_id):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    items = [
        {'row_id': 1, 'partition_id': 2, 'prediction': 0.42, 'prediction_values': []},
        {'row_id': 2, 'partition_id': 2, 'prediction': 0.22, 'prediction_values': []},
        {'row_id': 3, 'partition_id': 2, 'prediction': 0.82, 'prediction_values': []},
    ]
    response = {'data': items, 'next': None, 'prev': None}
    responses.add(
        responses.GET,
        url,
        body=json.dumps(response),
        status=200,
        content_type='application/json',
    )

    obj = TrainingPredictions.get(project_id, prediction_id)
    assert len(list(obj.iterate_rows(batch_size=3))) == 3


@responses.activate
def test_iterate_rows__one_prediction__ok(project_id, prediction_id):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    items = [
        {'row_id': 1, 'partition_id': 2, 'prediction': 0.42, 'prediction_values': []},
    ]
    response = {'data': items, 'next': '{}?offset=100&limit=100'.format(url)}
    responses.add(
        responses.GET,
        url,
        body=json.dumps(response),
        status=200,
        content_type='application/json',
    )

    obj = TrainingPredictions.get(project_id, prediction_id)

    assert list(obj.iterate_rows()) == [
        TrainingPredictionsRow(
            row_id=1,
            partition_id=2,
            prediction=0.42,
            prediction_values=[],
            timestamp=None,
            forecast_point=None,
            forecast_distance=None,
            series_id=None,
            prediction_explanations=None,
            shap_metadata=None,
        ),
    ]
    assert len(responses.calls) == 1


@responses.activate
def test_iterate_rows__one_prediction_with_shap_exp(project_id, prediction_id):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    items = [
        {
            'row_id': 1,
            'partition_id': 2,
            'prediction': 0.42,
            'prediction_values': [],
            'prediction_explanations': [
                {
                    'feature': "first_name",
                    'feature_value': "some_value",
                    'strength': -0.07,
                    'label': 1.0,
                },
                {
                    'feature': "second_name",
                    'feature_value': "0",
                    'strength': -0.05,
                    'label': 1.0,
                },
            ],
            'shap_metadata': {
                'shap_base_value': 0.39,
                'shap_remaining_total': 0.03,
                'warnings': {'mismatch_row_count': 1, 'max_normalized_mismatch': 0.01},
            },

        },
    ]
    response = {'data': items, 'next': '{}?offset=100&limit=100'.format(url)}
    responses.add(
        responses.GET, url, body=json.dumps(response), status=200, content_type='application/json'
    )

    training_predictions = TrainingPredictions.get(project_id, prediction_id)

    data = list(training_predictions.iterate_rows())
    assert len(data) == 1
    training_prediction = data[0]
    assert training_prediction.row_id == 1
    assert training_prediction.partition_id == 2
    assert training_prediction.prediction == 0.42
    assert training_prediction.prediction_values == []
    assert training_prediction.timestamp is None
    assert training_prediction.forecast_point is None
    assert training_prediction.forecast_distance is None
    assert training_prediction.series_id is None
    assert training_prediction.series_id is None
    assert training_prediction.prediction_explanations == [
        {
            'feature': "first_name",
            'feature_value': "some_value",
            'strength': -0.07,
            'label': 1.0,
        },
        {
            'feature': "second_name",
            'feature_value': "0",
            'strength': -0.05,
            'label': 1.0,
        }
    ]
    assert training_prediction.shap_metadata == {
        'shap_base_value': 0.39,
        'shap_remaining_total': 0.03,
        'warnings': {'mismatch_row_count': 1, 'max_normalized_mismatch': 0.01}
    }
    assert len(responses.calls) == 1


@responses.activate
def test_iterate_rows__two_predictions_with_limit_equals_one__three_requests_were_made(
        project_id, prediction_id,
):
    items = [
        {'row_id': 1, 'partition_id': 2, 'prediction': 0.42, 'prediction_values': []},
        {'row_id': 2, 'partition_id': 2, 'prediction': 0.22, 'prediction_values': []},
    ]
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    body = {'data': items[:1], 'next': '{}?offset=1&limit=1'.format(url)}
    responses.add(
        responses.GET,
        url,
        body=json.dumps(body),
        status=200,
        content_type='application/json',
    )

    obj = TrainingPredictions.get(project_id, prediction_id)

    it = obj.iterate_rows(batch_size=1)

    assert next(it, None) == TrainingPredictionsRow(
        row_id=1,
        partition_id=2,
        prediction=0.42,
        prediction_values=[],
        timestamp=None,
        forecast_point=None,
        forecast_distance=None,
        series_id=None,
        prediction_explanations=None,
        shap_metadata=None,
    )
    assert len(responses.calls) == 1

    body2 = {'data': items[1:], 'next': '{}?offset=2&limit=1'.format(url)}
    responses.reset()
    responses.add(
        responses.GET,
        url,
        body=json.dumps(body2),
        status=200,
        content_type='application/json',
        match_querystring=False,
    )

    assert next(it, None) == TrainingPredictionsRow(
        row_id=2,
        partition_id=2,
        prediction=0.22,
        prediction_values=[],
        timestamp=None,
        forecast_point=None,
        forecast_distance=None,
        series_id=None,
        prediction_explanations=None,
        shap_metadata=None,
    )
    assert len(responses.calls) == 1

    body = {'data': [], 'next': '{}?offset=3&limit=1'.format(url)}
    responses.reset()
    responses.add(
        responses.GET,
        url,
        body=json.dumps(body),
        status=200,
        content_type='application/json',
    )

    assert next(it, None) is None
    assert len(responses.calls) == 1


@responses.activate
@pytest.mark.parametrize(['serializer', 'response'], [
    pytest.param('json', {
        'body': lambda url: json.dumps({'data': [
            {'row_id': 1, 'partition_id': 2, 'prediction': 0.42, 'prediction_values': []},
            {'row_id': 2, 'partition_id': 2, 'prediction': 0.22, 'prediction_values': []},
        ], 'next': '{}?limit=20'.format(url)}),
        'content_type': 'application/json',
    }, id='json'),
    pytest.param('csv', {
        'body': lambda url: _write_csv([
            ['row_id', 'partition_id', 'prediction'],
            [1, 2, 0.42],
            [2, 2, 0.22],
        ]),
        'content_type': 'text/csv',
        'stream': True,
    }, id='csv')
])
def test_get_all_as_dataframe__regression_two_predictions__ok(
        project_id, prediction_id,
        project_url, project_with_target_json,
        serializer, response,
):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    response['body'] = response['body'](url)
    responses.add(
        responses.GET,
        url,
        status=200,
        **response
    )
    responses.add(
        responses.GET,
        project_url,
        body=project_with_target_json,
        status=200,
        content_type='application/json'
    )

    obj = TrainingPredictions.get(project_id, prediction_id)
    data_frame = obj.get_all_as_dataframe(serializer=serializer)

    assert isinstance(data_frame, pd.DataFrame)
    assert data_frame.shape == (2, 3)
    assert list(data_frame.columns) == [
        'row_id',
        'partition_id',
        'prediction',
    ]


@responses.activate
@pytest.mark.parametrize(['serializer', 'response'], [
    pytest.param('json', {
        'body': lambda url: json.dumps({'data': [
            {
                'row_id': 1,
                'partition_id': '1960-01-01T00:00:00.000000Z',
                'prediction': 0.42,
                'prediction_values': [],
                'timestamp': '1960-01-01T00:00:00.000000Z',
                'forecast_point': '1959-12-01T00:00:00.000000Z',
                'forecast_distance': 1,
                'series_id': "1",
            },
            {
                'row_id': 2,
                'partition_id': '1960-01-01T00:00:00.000000Z',
                'prediction': 0.22,
                'prediction_values': [],
                'timestamp': '1959-11-01T00:00:00.000000Z',
                'forecast_point': '1959-12-01T00:00:00.000000Z',
                'forecast_distance': 2,
                'series_id': "1",
            },
        ], 'next': '{}?limit=20'.format(url)}),
        'content_type': 'application/json',
    }, id='json'),
    pytest.param('csv', {
        'body': lambda url: _write_csv([
            [
                'row_id', 'partition_id', 'prediction', 'timestamp',
                'forecast_point', 'forecast_distance', 'series_id'
            ],
            [
                1,
                '1960-01-01T00:00:00.000000Z',
                0.42,
                '1960-01-01T00:00:00.000000Z',
                '1959-12-01T00:00:00.000000Z',
                1,
                "1",
            ],
            [
                2,
                '1960-01-01T00:00:00.000000Z',
                0.22,
                '1959-11-01T00:00:00.000000Z',
                '1959-12-01T00:00:00.000000Z',
                2,
                "1",
            ],
        ]),
        'content_type': 'text/csv',
        'stream': True,
    }, id='csv'),
    pytest.param('json', {
        'body': lambda url: json.dumps({'data': [
            {
                'row_id': 1,
                'partition_id': '1960-01-01T00:00:00.000000Z',
                'prediction': 0.42,
                'prediction_values': [],
                'timestamp': '1960-01-01T00:00:00.000000Z',
                'forecast_point': '1959-12-01T00:00:00.000000Z',
                'forecast_distance': 1,
                'series_id': 1,
            },
            {
                'row_id': 2,
                'partition_id': '1960-01-01T00:00:00.000000Z',
                'prediction': 0.22,
                'prediction_values': [],
                'timestamp': '1959-11-01T00:00:00.000000Z',
                'forecast_point': '1959-12-01T00:00:00.000000Z',
                'forecast_distance': 2,
                'series_id': 1,
            },
        ], 'next': '{}?limit=20'.format(url)}),
        'content_type': 'application/json',
    }, id='json'),
    pytest.param('csv', {
        'body': lambda url: _write_csv([
            [
                'row_id', 'partition_id', 'prediction', 'timestamp',
                'forecast_point', 'forecast_distance', 'series_id'
            ],
            [
                1,
                '1960-01-01T00:00:00.000000Z',
                0.42,
                '1960-01-01T00:00:00.000000Z',
                '1959-12-01T00:00:00.000000Z',
                1,
                1,
            ],
            [
                2,
                '1960-01-01T00:00:00.000000Z',
                0.22,
                '1959-11-01T00:00:00.000000Z',
                '1959-12-01T00:00:00.000000Z',
                2,
                1,
            ],
        ]),
        'content_type': 'text/csv',
        'stream': True,
    }, id='csv')
])
def test_get_all_as_dataframe__timeseries_predictions__ok(
        project_id, prediction_id,
        project_url, project_with_target_json,
        serializer, response,
):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    response['body'] = response['body'](url)
    responses.add(
        responses.GET,
        url,
        status=200,
        **response
    )
    project_data = json.loads(project_with_target_json)
    project_data['partition']['use_time_series'] = True
    responses.add(
        responses.GET,
        project_url,
        body=json.dumps(project_data),
        status=200,
        content_type='application/json'
    )

    obj = TrainingPredictions.get(project_id, prediction_id)
    data_frame = obj.get_all_as_dataframe(serializer=serializer)

    assert isinstance(data_frame, pd.DataFrame)
    assert data_frame.shape == (2, 7)
    assert set(data_frame.columns) == {
        'row_id', 'partition_id', 'prediction', 'timestamp',
        'forecast_point', 'forecast_distance', 'series_id'
    }

    assert len(data_frame) == 2
    assert None not in data_frame.timestamp.unique()
    assert None not in data_frame.forecast_point.unique()
    assert all(isinstance(value, int) for value in data_frame.forecast_distance)


@responses.activate
@pytest.mark.parametrize(['serializer', 'response'], [
    pytest.param('json', {
        'body': lambda url: json.dumps({'data': [
            {'row_id': 1, 'partition_id': 2, 'prediction': 0.42, 'prediction_values': []},
            {'row_id': 2, 'partition_id': 2, 'prediction': 0.22, 'prediction_values': []},
        ], 'next': '{}?limit=20'.format(url)}),
        'content_type': 'application/json',
    }, id='json'),
    pytest.param('csv', {
        'body': lambda url: _write_csv([
            ['row_id', 'partition_id', 'prediction'],
            [1, 2, 0.42],
            [2, 2, 0.22],
        ]),
        'content_type': 'text/csv',
        'stream': True,
    }, id='csv')
])
def test_download_to_csv__two_regression_predictions__ok(
        project_id, prediction_id,
        project_url, project_with_target_json,
        serializer, response,
):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    response['body'] = response['body'](url)
    responses.add(
        responses.GET,
        url,
        status=200,
        **response
    )
    responses.add(
        responses.GET,
        project_url,
        body=project_with_target_json,
        status=200,
        content_type='application/json'
    )

    obj = TrainingPredictions.get(project_id, prediction_id)
    with tempfile.NamedTemporaryFile(mode='r') as a_file:
        obj.download_to_csv(a_file.name, serializer=serializer)
        data_frame = pd.read_csv(a_file.name)

    assert data_frame.shape == (2, 3)
    assert set(data_frame.columns) == {'row_id', 'partition_id', 'prediction'}


@responses.activate
def test_iterate_rows__one_classification_prediction__ok(project_id, prediction_id):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    items = [
        {
            'row_id': 1,
            'partition_id': 2,
            'prediction': 0.42,
            'prediction_values': [
                {'label': 'good', 'value': 0.8},
                {'label': 'bad', 'value': 0.2},
            ],
        },
    ]
    response = {'data': items, 'next': '{}?offset=100&limit=100'.format(url)}
    responses.add(
        responses.GET,
        url,
        body=json.dumps(response),
        status=200,
        content_type='application/json',
    )

    obj = TrainingPredictions.get(project_id, prediction_id)

    assert list(obj.iterate_rows()) == [
        TrainingPredictionsRow(
            row_id=1,
            partition_id=2,
            prediction=0.42,
            prediction_values=[
                {'label': 'good', 'value': 0.8},
                {'label': 'bad', 'value': 0.2},
            ],
            timestamp=None,
            forecast_point=None,
            forecast_distance=None,
            series_id=None,
            prediction_explanations=None,
            shap_metadata=None,
        ),
    ]
    assert len(responses.calls) == 1


@responses.activate
@pytest.mark.parametrize(['serializer', 'response'], [
    pytest.param('json', {
        'body': lambda url: json.dumps({'data': [{
            'row_id': 1,
            'partition_id': 2,
            'prediction': 0.42,
            'prediction_values': [
                {'label': 'good', 'value': 0.8},
                {'label': 'bad', 'value': 0.2},
            ],
        }], 'next': '{}?limit=20'.format(url)}),
        'content_type': 'application/json',
    }, id='json'),
    pytest.param('csv', {
        'body': lambda url: _write_csv([
            ['row_id', 'partition_id', 'prediction', 'foo_good', 'foo_bad'],
            [1, 2, 0.42, 0.8, 0.2],
        ]),
        'content_type': 'text/csv',
        'stream': True,
    }, id='csv')
])
def test_get_all_as_dataframe__one_classification_prediction__ok(
        project_id, prediction_id,
        project_url, project_with_target_json,
        serializer, response,
):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    response['body'] = response['body'](url)
    responses.add(
        responses.GET,
        url,
        status=200,
        **response
    )

    responses.add(
        responses.GET,
        project_url,
        body=project_with_target_json,
        status=200,
        content_type='application/json'
    )

    obj = TrainingPredictions.get(project_id, prediction_id)
    data_frame = obj.get_all_as_dataframe(
        class_prefix='foo_',
        serializer=serializer,
    )

    assert isinstance(data_frame, pd.DataFrame)
    assert data_frame.shape == (1, 5)
    assert set(data_frame.columns) == {
        'row_id', 'partition_id', 'prediction', 'foo_good', 'foo_bad'
    }


@responses.activate
@pytest.mark.parametrize(['serializer', 'response'], [
    pytest.param('json', {
        'body': lambda url: json.dumps({'data': [{
            'row_id': 1,
            'partition_id': 2,
            'prediction': 0.42,
            'prediction_values': [
                {'label': 'good', 'value': 0.3},
                {'label': 'bad', 'value': 0.2},
                {'label': 'ugly', 'value': 0.2},
                {'label': 'none', 'value': 0.0},
            ],
        }], 'next': '{}?limit=20'.format(url)}),
        'content_type': 'application/json',
    }, id='json'),
    pytest.param('csv', {
        'body': lambda url: _write_csv([
            [
                'row_id', 'partition_id', 'prediction', 'class_good', 'class_bad',
                'class_ugly', 'class_none',
            ],
            [1, 2, 0.42, 0.3, 0.2, 0.2, 0.0],
        ]),
        'content_type': 'text/csv',
        'stream': True,
    }, id='csv')
])
def test_download_to_csv__multiclass_project__prediction__ok(
        project_id, prediction_id,
        project_url, project_with_target_json,
        serializer, response,
):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id,
        prediction_id,
    )
    response['body'] = response['body'](url)
    responses.add(
        responses.GET,
        url,
        status=200,
        **response
    )

    responses.add(
        responses.GET,
        project_url,
        body=project_with_target_json,
        status=200,
        content_type='application/json'
    )

    obj = TrainingPredictions.get(project_id, prediction_id)
    with tempfile.NamedTemporaryFile(mode='r') as a_file:
        obj.download_to_csv(a_file.name, serializer=serializer)
        data_frame = pd.read_csv(a_file.name)

    assert data_frame.shape == (1, 7)
    assert set(data_frame.columns) == {
        'row_id', 'partition_id', 'prediction', 'class_good', 'class_bad',
        'class_ugly', 'class_none',
    }


@responses.activate
@pytest.mark.parametrize(['serializer', 'response'], [
    pytest.param('json', {
        'body': lambda url: json.dumps({'data': [{
            'row_id': 0,
            'partition_id': "0.0",
            'prediction': 0.0,
            'prediction_threshold': 0.5,
            'prediction_values': [
                {'label': 1.0, 'value': 0.311178247734139},
                {'label': 0.0, 'value': 0.6888217522658611},
            ],
            'prediction_explanations': [
                {
                    'featureValue': '5',
                    'strength': -0.07643818895488697,
                    'feature': 'number_diagnoses',
                    'label': 1.0,
                },
            ],
            'shap_metadata': {
                "shap_base_value": 0.39640626311302185,
                "shap_remaining_total": 0.06644934718939229,
                "warnings": None,
            },
        }], 'next': '{}?limit=20'.format(url)}),
        'content_type': 'application/json',
    }, id='json'),
    pytest.param('csv', {
        'body': lambda url: _write_csv([
            [
                'row_id', 'partition_id', 'prediction', 'class_1.0', 'class_0.0', 'explained_class',
                'Explanation_1_feature_name', 'Explanation_1_feature_value',
                'Explanation_1_strength', 'shap_remaining_total', 'shap_base_value'
            ],
            [
                0, 0.0, 0.0, 0.311178247734139, 0.6888217522658611, 1.0, 'number_diagnoses', 5,
                -0.07643818895488697, 'number_inpatient', 0, -0.05988679910481005,
                0.06644934718939229, 0.396406263113
            ]
        ]),
        'content_type': 'text/csv',
        'stream': True,
    }, id='csv')
])
def test_get_all_as_dataframe_one_classification_prediction_with_explanations(
        project_id, prediction_id, project_url, project_with_target_json, serializer, response
):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id, prediction_id
    )
    response['body'] = response['body'](url)
    responses.add(responses.GET, url, status=200, **response)

    responses.add(
        responses.GET,
        project_url,
        body=project_with_target_json,
        status=200,
        content_type='application/json'
    )

    obj = TrainingPredictions.get(project_id, prediction_id)
    data_frame = obj.get_all_as_dataframe(serializer=serializer)

    assert isinstance(data_frame, pd.DataFrame)
    assert data_frame.shape == (1, 11)
    assert set(data_frame.columns) == {
        'row_id', 'partition_id', 'prediction', 'class_1.0', 'class_0.0', 'explained_class',
        'shap_remaining_total', 'shap_base_value',
        'Explanation_1_feature_name', 'Explanation_1_feature_value', 'Explanation_1_strength',
    }


@responses.activate
@pytest.mark.parametrize(['serializer', 'response'], [
    pytest.param('json', {
        'body': lambda url: json.dumps({'data': [
            {
                'row_id': 0,
                'partition_id': '0.0',
                'prediction': 71700.0,
                'prediction_values': [
                    {'value': 71700.0, 'label': 'SalePrice'}
                ],
                'shap_metadata': {
                    'shap_base_value': 31108.28125,
                    'shap_remaining_total': -4045.506156357904,
                    'warnings': None,
                },
                'prediction_explanations': [
                    {
                        'feature_value': '172632',
                        'strength': 9021.030323424387,
                        'feature': 'MachineID',
                        'label': 'SalePrice',
                    },
                    {
                        'feature_value': '3855',
                        'strength': -4883.805345350174,
                        'feature': 'ModelID',
                        'label': 'SalePrice',
                    }
                ],
            },
            {
                'row_id': 1,
                'partition_id': '0.0',
                'prediction': 12275.0,
                'prediction_values': [
                    {'value': 12275.0, 'label': 'SalePrice'}
                ],
                'shap_metadata': {
                    'shap_base_value': 31108.28125,
                    'shap_remaining_total': -3350.9689739794676,
                    'warnings': None,
                },
                'prediction_explanations': [
                    {
                        'feature_value': '1986',
                        'strength': -6307.912477762799,
                        'feature': 'YearMade',
                        'label': 'SalePrice',
                    },
                    {
                        'feature_value': '1190700',
                        'strength': -3257.09113279142,
                        'feature': 'MachineID',
                        'label': 'SalePrice',
                    }
                ],
            },
        ], 'next': '{}?limit=20'.format(url)}),
        'content_type': 'application/json',
    }, id='json'),
    pytest.param('csv', {
        'body': lambda url: _write_csv([
            [
                'row_id', 'partition_id', 'prediction', 'Explanation_1_feature_name',
                'Explanation_1_feature_value', 'Explanation_1_strength',
                'Explanation_2_feature_name', 'Explanation_2_feature_value',
                'Explanation_2_strength', 'shap_remaining_total', 'shap_base_value'
            ],
            [
                0, 0.0, 71700.0, 'MachineID', '172632', 9021.030323424387, 'ModelID', '3855',
                -4883.805345350174, -4045.506156357904, 31108.28125
            ],
            [
                1, 0.0, 12275.0, 'YearMade', '1986', -6307.912477762799, 'MachineID', '1190700',
                -3257.09113279142, -3350.9689739794676, 31108.28125
            ],
        ]),
        'content_type': 'text/csv',
        'stream': True,
    }, id='csv')
])
def test_get_all_as_dataframe_regression_predictions_with_explanations(
        project_id, prediction_id, project_url, project_with_target_json, serializer, response
):
    url = 'https://host_name.com/projects/{}/trainingPredictions/{}/'.format(
        project_id, prediction_id
    )
    response['body'] = response['body'](url)
    responses.add(responses.GET, url, status=200, **response)
    responses.add(
        responses.GET,
        project_url,
        body=project_with_target_json,
        status=200,
        content_type='application/json',
    )

    predictions = TrainingPredictions.get(project_id, prediction_id)
    data_frame = predictions.get_all_as_dataframe(serializer=serializer)

    assert isinstance(data_frame, pd.DataFrame)
    assert data_frame.shape == (2, 11)
    assert list(data_frame.columns) == [
        'row_id', 'partition_id', 'prediction', 'Explanation_1_feature_name',
        'Explanation_1_feature_value', 'Explanation_1_strength', 'Explanation_2_feature_name',
        'Explanation_2_feature_value', 'Explanation_2_strength', 'shap_remaining_total',
        'shap_base_value'
    ]
