# -*- coding: utf-8 -*-
import json

import numpy as np
import pandas as pd
import pytest
import responses
import six
from trafaret import DataError

from datarobot._experimental.models.model import CombinedModel
from datarobot._experimental.models.project import Project
from tests.utils import add_response

_project_id = "projectId"
_target = "target"
_model_id = "modelId"
_base_url = "https://host_name.com/projects/{}/combinedModels/".format(_project_id)
_valid_combined_model_id = "validCombinedModelId"
_invalid_combined_model_id = "validCombinedModelId"
_valid_combined_model_url = "{0}{1}/".format(_base_url, _valid_combined_model_id)
_invalid_combined_model_url = "{0}{1}/".format(_base_url, _invalid_combined_model_id)
_combined_model_segments_dl_url = "{0}{1}/segments/download/".format(
    _base_url, _valid_combined_model_id
)
_segment_champion_url = "https://host_name.com/projects/{}/segmentChampion/".format(_project_id)


@pytest.fixture
def combined_model_record():
    return {
        "combinedModelId": _valid_combined_model_id,
        "projectId": _project_id,
        "segmentationTaskId": "segmentationTaskId",
    }


@pytest.fixture
def combined_model_list_response(combined_model_record):
    return {"count": 1, "next": None, "previous": None, "data": [combined_model_record]}


@pytest.fixture
def combined_model_get_json(combined_model_record):
    return json.dumps(combined_model_record)


@pytest.fixture
def invalid_json():
    return json.dumps({"message": "Not Found"})


@pytest.fixture
def combined_model_list_json(combined_model_list_response):
    return json.dumps(combined_model_list_response)


@pytest.fixture
def mock_segment_csv():
    row_indexes = [u"row_{}".format(i) for i in range(10)]
    column_names = [u"カテゴリー_{}".format(i) for i in range(5)]
    rs = np.random.RandomState(1234)
    df = pd.DataFrame(
        data=rs.uniform(-1, 1, (len(row_indexes), len(column_names))),
        columns=column_names,
        index=row_indexes,
    )
    return df.to_csv(encoding="utf-8")


@responses.activate
def test_combined_model_list(combined_model_list_json):
    responses.add(
        responses.GET,
        _base_url,
        body=combined_model_list_json,
        status=200,
        content_type="application/json",
    )
    project = Project(id=_project_id)
    combined_models = project.get_combined_models()
    assert len(combined_models) == 1
    combined_model = combined_models[0]

    assert isinstance(combined_model, CombinedModel)
    assert combined_model.id == _valid_combined_model_id
    assert combined_model.project_id == _project_id


@responses.activate
def test_combined_model_get_valid(combined_model_get_json):
    responses.add(
        responses.GET,
        _valid_combined_model_url,
        body=combined_model_get_json,
        status=200,
        content_type="application/json",
    )
    combined_model = CombinedModel.get(_project_id, _valid_combined_model_id)

    assert isinstance(combined_model, CombinedModel)
    assert combined_model.id == _valid_combined_model_id
    assert combined_model.project_id == _project_id


@responses.activate
def test_combined_model_invalid_get(invalid_json):
    responses.add(responses.GET, _invalid_combined_model_url, body=invalid_json)
    with pytest.raises(DataError):
        CombinedModel.get(_project_id, _invalid_combined_model_id)


@responses.activate
def test_combined_model_get_segments_as_dataframe(mock_segment_csv):
    add_response(_combined_model_segments_dl_url, mock_segment_csv, content_type="text/csv")
    combined_model = CombinedModel(_valid_combined_model_id, _project_id)
    df = combined_model.get_segments_as_dataframe()

    assert isinstance(df, pd.DataFrame)
    assert df.shape == (10, 5)


@responses.activate
def test_combined_model_get_segments_as_csv(mock_segment_csv):
    add_response(_combined_model_segments_dl_url, mock_segment_csv, content_type="text/csv")
    combined_model = CombinedModel(_valid_combined_model_id, _project_id)
    file_buffer = six.StringIO()
    combined_model.get_segments_as_csv(file_buffer)
    file_buffer.seek(0)

    # Assert the values contained in the output reasonably match the inputs
    actual_values = pd.read_csv(six.StringIO(file_buffer.read())).values
    expected_values = pd.read_csv(six.StringIO(mock_segment_csv)).values[:, 1:]
    np.testing.assert_almost_equal(actual_values, expected_values)


@responses.activate
def test_combined_model_valid_set_segment_champion(invalid_json):
    responses.add(
        responses.PUT,
        _segment_champion_url,
        body=json.dumps({"combinedModelId": str(_model_id)}),
        status=200,
        content_type="application/json",
    )

    combined_model = CombinedModel.set_segment_champion(_project_id, _model_id)
    assert combined_model == "modelId"
