import json
import os
from datetime import datetime, timezone

import arthurai.core.util as core_util
import numpy as np
import pandas as pd
import pytest
import responses

from arthurai.common.constants import ValueType, Stage, InputType, OutputType, ImageResponseType, ImageContentType, \
    ModelStatus, ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE
from arthurai.common.exceptions import UserValueError, MethodNotApplicableError, ArthurInternalError
from arthurai.core.attributes import AttributeCategory, AttributeBin, ArthurAttribute
from arthurai.core.models import ArthurModel, ArthurModelGroup
from http import HTTPStatus
from mock import patch, mock_open, MagicMock
from tests.helpers import assert_attributes_equal, mock_post, mock_get, mock_patch
from tests.fixtures.models import regression as reg, multiclass as mclass, nlp
from tests.fixtures.models.regression import initial_batch_model as reg_initial_batch_model, \
    unsaved_batch_model as reg_unsaved_batch_model, streaming_model
from tests.fixtures.models.multiclass import initial_biclass_model, SAMPLE_DATAFRAME_NO_NEG_PRED, \
    MODEL_DATA_WITH_ATTRIBUTES_IMPLICIT_NEG
from tests.fixtures.models.nlp import initial_nlp_classification_model, initial_nlp_regression_model
from tests.fixtures.mocks import client
from tests.test_request_models import load_json, load_json_string
from tests.test_request_models.fixtures import model_response_json_strings


MODEL_RESPONSE_JSON = [
    load_json(
        os.path.join(os.path.dirname(__file__), 'test_request_models', 'resources', 'model_with_categories.json')),
    load_json(
        os.path.join(os.path.dirname(__file__), 'test_request_models', 'resources', 'model_with_attributes.json'))
]

MODEL_GROUP_ID = "01234567-89ab-cdef-0123-456789abcdef"
MODEL_GROUP_NAME = "Test Model Group"
MODEL_GROUP_DESCRIPTION = "Test Model Group Description"
MODEL_GROUP_ARCHIVED = False
MODEL_GROUP_CREATED_AT = "2022-01-20T23:22:18.185267Z"
MODEL_GROUP_UPDATED_AT = "2022-01-20T23:22:32.914068Z"
MODEL_GROUP_JSON = {
    "id": MODEL_GROUP_ID,
    "name": MODEL_GROUP_NAME,
    "description": MODEL_GROUP_DESCRIPTION,
    "archived": MODEL_GROUP_ARCHIVED,
    "created_at": MODEL_GROUP_CREATED_AT,
    "updated_at": MODEL_GROUP_UPDATED_AT,
    "versions": MODEL_RESPONSE_JSON
}


@pytest.mark.parametrize('model_filename,expected_string', [
    ('model_with_attributes.json', "Tabular Regression Model\n"
                                   "\tFICO Score:3456\n"
                                   "\tStage: Production\n"
                                   "\tAttributes: ['PIPELINE_INPUT[2]:attr_3']"),
    ('model_multiclass_with_bias.json', "Tabular Multiclass Model\n"
                                        "\ttest model 1:1113\n"
                                        "\tStage: Production\n"
                                        "\tAttributes: ['NON_INPUT_DATA[0]:bias', 'PREDICTED_VALUE[0]:attr_1', 'PREDICTED_VALUE[1]:attr_2', 'PREDICTED_VALUE[2]:attr_3']")
])
def test__str__(model_filename, expected_string):
    model_json = load_json_string(os.path.join(os.path.dirname(__file__), 'test_request_models', 'resources', model_filename))
    model = ArthurModel.from_json(model_json)
    assert str(model) == expected_string


@responses.activate
def test_send_parquet_file():
    input_dict = {
        "attr1": np.Inf,
        "attr2": "string1",
        "attr3": 3.44,
        "attr4": True,
        "attr5": np.nan,
        "attr6": -np.inf,
        "attr7": "string2",
        "attr8": None,
        "attr9": "",
        "attr10": np.inf,
    }
    expected_dict = {
        "attr1": None,
        "attr2": "string1",
        "attr3": 3.44,
        "attr4": True,
        "attr5": None,
        "attr6": None,
        "attr7": "string2",
        "attr8": None,
        "attr9": "",
        "attr10": None,
    }
    output_dict = ArthurModel._replace_nans_and_infinities_in_dict(input_dict)
    assert expected_dict == output_dict


def test_replace_nans_none_type():
    output_dict = ArthurModel._replace_nans_and_infinities_in_dict(None)
    assert output_dict is None


def test_standardize_pd_obj_for_nans():
    input_df = pd.DataFrame(
        [
            [1, 2, 3., 4, 5, 6, np.inf, np.nan, 1, 1],
            [7, 8, 9., 1, 2, 3, np.nan, -np.inf, 2, 2],
            [4, 5.5, 9., -np.inf, np.nan, np.Inf, np.nan, np.nan, np.nan, np.nan],
        ],
        columns = ["x0", "x1", "x2", "x3", "x4", "x5", "x6", "x7", "x8", "x9"]
    )
    expected_df = pd.DataFrame(
        [
            [1, 2, 3., 4, 5, 6, np.nan, np.nan, 1, 1],
            [7, 8, 9., 1, 2, 3, np.nan, np.nan, 2, 2],
            [4, 5.5, 9., np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
        ],
        columns = ["x0", "x1", "x2", "x3", "x4", "x5", "x6", "x7", "x8", "x9"]
    )
    expected_df.x3 = expected_df.x3.astype(pd.Int64Dtype())
    expected_df.x4 = expected_df.x4.astype(pd.Int64Dtype())
    expected_df.x5 = expected_df.x5.astype(pd.Int64Dtype())

    attributes = {
        "x3": ValueType.Integer,
        "x4": ValueType.Integer,
        "x5": ValueType.Integer,
        "x8": ValueType.Float,
    }

    assert expected_df.equals(
        core_util.standardize_pd_obj(input_df, dropna=False, replacedatetime=False, attributes=attributes)
    )
    with pytest.raises(ValueError):
        assert expected_df.equals(
            core_util.standardize_pd_obj(input_df, dropna=True, replacedatetime=False, attributes=attributes)
        )


def test_standardize_pd_obj_for_series():
    input_series = [
        pd.Series([1, 2, 3]),
        pd.Series([1., 2., 3.]),
        pd.Series([1, 2, np.Inf], name="x1"),
        pd.Series([1, 2, np.Inf], name="x1"),
        pd.Series([1, 2, np.Inf], name="x1"),
        pd.Series([np.nan, -np.inf, np.inf]),
    ]
    input_attributes = [
        {},
        None,
        {"x1": ValueType.Integer},
        {"x1": ValueType.Float},
        {"x2": ValueType.Integer},
        {}
    ]
    expected_undropped_series = [
        pd.Series([1, 2, 3]),
        pd.Series([1., 2., 3.]),
        pd.Series([1, 2, np.nan]).astype(pd.Int64Dtype()),
        pd.Series([1, 2, np.nan]),
        pd.Series([1, 2, np.nan]),
        pd.Series([np.nan, np.nan, np.nan]),
    ]
    expected_dropped_series = [
        pd.Series([1, 2, 3]),
        pd.Series([1., 2., 3.]),
        pd.Series([1, 2]).astype(pd.Int64Dtype()),
        pd.Series([1, 2.]),
        pd.Series([1, 2.]),
        pd.Series([]),
    ]

    for inp, inp_attributes, exp_undropped, exp_dropped in zip(input_series, input_attributes, expected_undropped_series, expected_dropped_series):
        print(inp_attributes)
        assert exp_undropped.equals(
            core_util.standardize_pd_obj(inp, dropna=False, replacedatetime=False, attributes=inp_attributes)
        )
        assert exp_dropped.equals(
            core_util.standardize_pd_obj(inp, dropna=True, replacedatetime=False, attributes=inp_attributes)
        )


def test_standardize_pd_obj_bad_input():
    for inp in ["string", "", 123, [1,2,3], {"key": "value"}]:
        with pytest.raises(TypeError):
            core_util.standardize_pd_obj(inp, dropna=True, replacedatetime=False)


def test_standardize_pd_obj_for_datetimes():
    unaware_timestamp = datetime(2021, 2, 8, 12, 10, 5)
    aware_timestamp = datetime(2021, 2, 8, 12, 10, 5, tzinfo=timezone.utc)
    timestamp_to_expected_output = (
        (unaware_timestamp, ValueError),
        (aware_timestamp, "2021-02-08T12:10:05+00:00"),
    )

    for timestamp, expected_output in timestamp_to_expected_output:
        if isinstance(expected_output, str):
            actual_output = core_util.standardize_pd_obj(pd.Series([timestamp]), dropna=False, replacedatetime=True)
            assert expected_output == actual_output.values[0]
        else:
            with pytest.raises(expected_output):
                core_util.standardize_pd_obj(pd.Series([timestamp]), dropna=False, replacedatetime=True)


def test_add_attribute_with_categories():
    model = ArthurModel.from_json(model_response_json_strings[2])
    model.add_attribute(
        name="test_categorical_1",
        value_type=ValueType.Integer,
        stage=Stage.ModelPipelineInput,
        categorical=True,
        categories=[1, 2, 3, 4],
    )
    attr = model.get_attribute(name="test_categorical_1")
    assert len(attr.categories) == 4
    assert isinstance(attr.categories[0], AttributeCategory)

    model.add_attribute(
        name="test_categorical_2",
        value_type=ValueType.Integer,
        stage=Stage.ModelPipelineInput,
        categorical=True,
        categories=[AttributeCategory(value="hello")],
    )
    attr = model.get_attribute(name="test_categorical_2")
    assert len(attr.categories) == 1
    assert isinstance(attr.categories[0], AttributeCategory)
    assert attr.categories[0].value == "hello"


def test_add_attribute_with_bins():
    model = ArthurModel.from_json(model_response_json_strings[2])
    model.add_attribute(
        name="test_bins_1",
        value_type=ValueType.Integer,
        stage=Stage.ModelPipelineInput,
        categorical=True,
        bins=[None, 10, 15, 20, None],
    )
    attr = model.get_attribute(name="test_bins_1")
    assert len(attr.bins) == 4
    assert isinstance(attr.bins[0], AttributeBin)
    assert attr.bins[0].continuous_start is None
    assert attr.bins[0].continuous_end == 10

    model.add_attribute(
        name="test_bins_2",
        value_type=ValueType.Integer,
        stage=Stage.ModelPipelineInput,
        categorical=True,
        bins=[
            AttributeBin(continuous_start=None, continuous_end=10),
            AttributeBin(continuous_start=10, continuous_end=15),
            AttributeBin(continuous_start=15, continuous_end=20),
        ],
    )
    attr = model.get_attribute(name="test_bins_2")
    assert len(attr.bins) == 3
    assert isinstance(attr.bins[0], AttributeBin)
    assert attr.bins[0].continuous_start is None
    assert attr.bins[0].continuous_end == 10


def test_add_nlp_attribute():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.NLP,
        output_type=OutputType.Multiclass,
    )
    ref_data = pd.DataFrame(
        {"input_text": ["hello", "hi", "what's up", "yo", "cool"],}
    )
    model.infer_schema(ref_data, stage=Stage.ModelPipelineInput)

    assert (
        model.get_attribute("input_text").value_type == ValueType.Unstructured_Text
    )
    assert model.get_attribute("input_text").categories is None


def test_float_categorical():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.Tabular,
        output_type=OutputType.Multiclass,
    )
    ref_data = pd.DataFrame(
        {"feature1": [float(1.0), float(1.0), float(-1.0), float(-1.0), float(1.0)], }
    )
    model.infer_schema(ref_data, stage=Stage.ModelPipelineInput)

    attribute = model.get_attribute("feature1")

    assert attribute.value_type == ValueType.Float
    assert attribute.categorical
    category_values = sorted([a.value for a in attribute.categories])
    assert category_values == ["-1.0", "1.0"]


def test_all_null_column():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.NLP,
        output_type=OutputType.Multiclass,
    )
    df = pd.DataFrame({'col1': [1], 'col2': [float("nan")]})
    model.infer_schema(df, stage=Stage.ModelPipelineInput)

    expected_attributes = [
        ArthurAttribute(name='col1', value_type=ValueType.Integer, stage=Stage.ModelPipelineInput, position=0,
                        categorical=True, categories=[AttributeCategory(value='1', label=None)], is_unique=False),
        ArthurAttribute(name='col2', value_type=ValueType.Float, stage=Stage.ModelPipelineInput, position=1,
                        categorical=False, is_unique=False)]

    assert model.attributes == expected_attributes


def test_nonzero_index_column():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.NLP,
        output_type=OutputType.Multiclass,
    )
    df = pd.DataFrame({'col1': [1], 'col2': [float("nan")]})
    df.index = [5]
    model.infer_schema(df, stage=Stage.ModelPipelineInput)

    expected_attributes = [
        ArthurAttribute(name='col1', value_type=ValueType.Integer, stage=Stage.ModelPipelineInput, position=0,
                        categorical=True, categories=[AttributeCategory(value='1', label=None)], is_unique=False),
        ArthurAttribute(name='col2', value_type=ValueType.Float, stage=Stage.ModelPipelineInput, position=1,
                        categorical=False, is_unique=False)]

    assert model.attributes == expected_attributes


def test_infer_numerical_string():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.Tabular,
        output_type=OutputType.Multiclass,
    )
    df = pd.DataFrame({'col1': ["0", "1", "0", "1", "1", "0"], 'col2': ["0.0", "1.0", "0.0", "1.0", "1.0", "0.0"]})
    model.infer_schema(df, stage=Stage.ModelPipelineInput)

    expected_attributes = [
        ArthurAttribute(name='col1', value_type=ValueType.String, stage=Stage.ModelPipelineInput, position=0,
                        categorical=True, categories=[AttributeCategory(value='0', label=None),
                                                      AttributeCategory(value='1', label=None)], is_unique=False),
        ArthurAttribute(name='col2', value_type=ValueType.String, stage=Stage.ModelPipelineInput, position=1,
                        categorical=True, categories=[AttributeCategory(value='0.0', label=None),
                                                      AttributeCategory(value='1.0', label=None)], is_unique=False)
    ]

    assert_attributes_equal(expected_attributes, model.attributes)


def test_infer_string_categorical_dtype():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.Tabular,
        output_type=OutputType.Multiclass
    )

    df = pd.DataFrame({"col1": list("abca"), "col2": list("bccd")}, dtype="category")
    model.infer_schema(df, stage=Stage.ModelPipelineInput)

    expected_attributes = [
        ArthurAttribute(name='col1', value_type=ValueType.String, stage=Stage.ModelPipelineInput, position=0,
                        categorical=True, categories=[AttributeCategory(value='a', label=None),
                                                      AttributeCategory(value='b', label=None),
                                                      AttributeCategory(value='c', label=None)], is_unique=False),
        ArthurAttribute(name='col2', value_type=ValueType.String, stage=Stage.ModelPipelineInput, position=1,
                        categorical=True, categories=[AttributeCategory(value='d', label=None),
                                                      AttributeCategory(value='b', label=None),
                                                      AttributeCategory(value='c', label=None)], is_unique=False)
    ]
    assert_attributes_equal(expected_attributes, model.attributes)


def test_infer_int_categorical_dtype():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.Tabular,
        output_type=OutputType.Multiclass
    )

    df = pd.DataFrame({"col1": [1, 0, 0, 1], "col2": [1, 2, 1, 0]}, dtype="category")
    model.infer_schema(df, stage=Stage.ModelPipelineInput)

    expected_attributes = [
        ArthurAttribute(name='col1', value_type=ValueType.Integer, stage=Stage.ModelPipelineInput, position=0,
                        categorical=True, categories=[AttributeCategory(value='0', label=None),
                                                      AttributeCategory(value='1', label=None)], is_unique=False),
        ArthurAttribute(name='col2', value_type=ValueType.Integer, stage=Stage.ModelPipelineInput, position=1,
                        categorical=True, categories=[AttributeCategory(value='0', label=None),
                                                      AttributeCategory(value='1', label=None),
                                                      AttributeCategory(value='2', label=None)], is_unique=False)
    ]
    assert_attributes_equal(expected_attributes, model.attributes)


def test_set_categorical_value_labels():
    model = ArthurModel.from_json(model_response_json_strings[2])
    model.add_attribute(
        name="test_categorical_1",
        value_type=ValueType.Integer,
        stage=Stage.ModelPipelineInput,
        categorical=True,
        categories=[1, 2],
    )
    expected_categories = {1: "Male", 2: "Female"}
    model.set_attribute_labels("test_categorical_1", labels=expected_categories)
    cats = model.get_attribute("test_categorical_1").categories
    for category in cats:
        assert int(category.value) in expected_categories
        assert expected_categories[int(category.value)] == category.label


@responses.activate
def test_get_image(client):
    model = ArthurModel(client=client.client,
                        partner_model_id="test",
                        input_type=InputType.Image,
                        output_type=OutputType.Multiclass)
    model.id = '1234567890abcdef'
    image_id = '0a1b2c3d4e5f9687'
    path = '.'
    type = ImageResponseType.RawImage
    file_ext = '.png'
    response_content = 'content'.encode()
    expected_image_file = f"{path}/{type}_{image_id}{file_ext}"

    mock_get(f"/api/v3/models/{model.id}/inferences/images/{image_id}?type={type}", response_content,
             headers={'Content-Type': ImageContentType.Png})
    open_mock = mock_open()
    with patch("arthurai.core.models.open", open_mock, create=True):
        resulting_file = model.get_image(image_id, path, type=type)

    assert resulting_file == expected_image_file
    open_mock.assert_called_with(expected_image_file, 'wb')
    open_mock.return_value.write.assert_called_once_with(response_content)


def test_add_object_detection_output_attributes():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.Image,
        output_type=OutputType.Multiclass,
    )

    # cannot call without proper types
    with pytest.raises(MethodNotApplicableError):
        model.add_object_detection_output_attributes('pred', 'gt', ['class_1', 'class_2'])
    model.input_type = InputType.NLP
    model.output_type = OutputType.ObjectDetection
    with pytest.raises(MethodNotApplicableError):
        model.add_object_detection_output_attributes('pred', 'gt', ['class_1', 'class_2'])

    # cannot call without classes
    model.output_type = OutputType.ObjectDetection
    model.input_type = InputType.Image
    with pytest.raises(UserValueError):
        model.add_object_detection_output_attributes('pred', 'gt', [])

    # names cannot match
    with pytest.raises(UserValueError):
        model.add_object_detection_output_attributes('pred', 'pred', ['class_1', 'class_2'])

    # happy path
    model.add_object_detection_output_attributes('pred', 'gt', ['class_1', 'class_2'])
    assert model.get_attribute('pred').value_type == ValueType.BoundingBox
    assert model.get_attribute('gt').value_type == ValueType.BoundingBox
    assert model.image_class_labels == ['class_1', 'class_2']


def test_build_tabular_model(reg_initial_batch_model):
    pred_gt_map = {reg.PRED: reg.GROUND_TRUTH}
    reg_initial_batch_model.build(data=reg.SAMPLE_DATAFRAME, pred_to_ground_truth_map=pred_gt_map,
                                  non_input_columns=[reg.FLOAT_NONINPUT])

    assert_attributes_equal(reg.MODEL_ATTRIBUTES, reg_initial_batch_model.attributes)
    pd.testing.assert_frame_equal(reg_initial_batch_model.reference_dataframe, reg.SAMPLE_DATAFRAME)


def test_build_tabular_class_model(initial_biclass_model):
    pred_gt_map = {mclass.DOG_PRED: mclass.DOG_GROUND_TRUTH, mclass.CAT_PRED: mclass.CAT_GROUND_TRUTH}
    initial_biclass_model.build(data=mclass.SAMPLE_DATAFRAME, pred_to_ground_truth_map=pred_gt_map,
                                non_input_columns=[mclass.FLOAT_NONINPUT], positive_predicted_attr=mclass.DOG_PRED)

    assert_attributes_equal(mclass.MODEL_ATTRIBUTES, initial_biclass_model.attributes)
    pd.testing.assert_frame_equal(initial_biclass_model.reference_dataframe, mclass.SAMPLE_DATAFRAME)


def test_build_tabular_class_model_implicit_negative_pred(initial_biclass_model):
    pred_gt_map = {mclass.DOG_PRED: mclass.DOG_GROUND_TRUTH, None: mclass.CAT_GROUND_TRUTH}
    initial_biclass_model.build(data=SAMPLE_DATAFRAME_NO_NEG_PRED,
                                pred_to_ground_truth_map=pred_gt_map, non_input_columns=[mclass.FLOAT_NONINPUT])
    assert_attributes_equal(mclass.MODEL_ATTRIBUTES_NO_NEG_PRED, initial_biclass_model.attributes)
    pd.testing.assert_frame_equal(initial_biclass_model.reference_dataframe, SAMPLE_DATAFRAME_NO_NEG_PRED)


def test_build_tabular_model_no_refset(reg_unsaved_batch_model):
    pred_gt_map = {reg.PRED: reg.GROUND_TRUTH}
    reg_unsaved_batch_model.build(data=reg.SAMPLE_DATAFRAME, pred_to_ground_truth_map=pred_gt_map,
                                  non_input_columns=[reg.FLOAT_NONINPUT], set_reference_data=False)

    assert_attributes_equal(reg.MODEL_ATTRIBUTES, reg_unsaved_batch_model.attributes)
    assert reg_unsaved_batch_model.reference_dataframe is None


def test_build_tabular_model_single_column_gt(initial_biclass_model):
    pred_class_map = {mclass.DOG_PRED: "0", mclass.CAT_PRED: "1"}
    initial_biclass_model.build(data=mclass.SINGLE_COLUMN_DATAFRAME, pred_to_ground_truth_map=pred_class_map,
                                ground_truth_column=mclass.GROUND_TRUTH, non_input_columns=[mclass.FLOAT_NONINPUT],
                                positive_predicted_attr=mclass.DOG_PRED)
    assert_attributes_equal(mclass.SINGLE_COL_MODEL_ATTRIBUTES, initial_biclass_model.attributes)
    pd.testing.assert_frame_equal(initial_biclass_model.reference_dataframe, mclass.SINGLE_COLUMN_DATAFRAME)


def test_build_nlp(initial_nlp_regression_model):
    initial_nlp_regression_model.build(nlp.NLP_DATAFRAME, pred_to_ground_truth_map={"pred": "gt"})

    assert initial_nlp_regression_model.get_attribute("input_text").value_type == ValueType.Unstructured_Text
    assert initial_nlp_regression_model.get_attribute("input_text").categories is None
    pd.testing.assert_frame_equal(initial_nlp_regression_model.reference_dataframe, nlp.NLP_DATAFRAME)


def test_build_nlp_sing_col(initial_nlp_classification_model):
    pred_class_map = {mclass.DOG_PRED: nlp.DOG_CLASS_STRING, mclass.CAT_PRED: nlp.CAT_CLASS_STRING}
    initial_nlp_classification_model.build(data=nlp.NLP_CLASSIFICATION_DATAFRAME,
                                           pred_to_ground_truth_map=pred_class_map,
                                           positive_predicted_attr=mclass.DOG_PRED,
                                           ground_truth_column=mclass.GROUND_TRUTH)
    assert_attributes_equal(nlp.CLASSIFICATION_MODEL_ATTRIBUTES, initial_nlp_classification_model.attributes)
    pd.testing.assert_frame_equal(initial_nlp_classification_model.reference_dataframe, nlp.NLP_CLASSIFICATION_DATAFRAME)


def test_build_image():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.Image,
        output_type=OutputType.Regression,
    )
    data = pd.DataFrame({'image': ["path1", "path2"], 'preds': [15.2, 31.6], 'gt': [12, 40]})
    model.build(data, pred_to_ground_truth_map={'preds': 'gt'})

    assert model.get_attribute("image").value_type == ValueType.Image
    assert model.get_attribute("image").categories is None
    pd.testing.assert_frame_equal(model.reference_dataframe, data)


def test_build_object_detection():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.Image,
        output_type=OutputType.ObjectDetection,
    )
    # not actually valid object data but that's okay
    data = pd.DataFrame({'image': ["path1", "path2"], 'preds': [0.8, 0.3], 'gt': [1, 0]})
    with pytest.raises(MethodNotApplicableError):
        model.build(data, pred_to_ground_truth_map={'preds': 'gt'})

def test_build_binary_classifier():
    model = ArthurModel(
        partner_model_id="test",
        client=None,
        input_type=InputType.Tabular,
        output_type=OutputType.Multiclass,
        classifier_threshold = 0.42
    )

    data = pd.DataFrame({'pred0': [0.1,0.9,0.2,0.8,0.3,0.7],
                        'pred1': [0.9,0.1,0.8,0.2,0.7,0.3],
                        'gt0': [1,1,1,0,0,0],
                        'gt1': [0,0,0,1,1,1]})
    
    model.build(data, pred_to_ground_truth_map = {'pred0': 'gt0', 'pred1': 'gt1'}, positive_predicted_attr= 'pred1')

    assert model.classifier_threshold == 0.42



@responses.activate
def test_save_model_with_refset(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data

    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': [],
        'status': ModelStatus.Pending
    }, status=201)

    with patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
               return_value=ModelStatus.Ready) as model_status_waiter_mock:
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    model_status_waiter_mock.assert_called_once_with([ModelStatus.Ready, ModelStatus.CreationFailed],
                                                     ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1
    assert reg_unsaved_batch_model.status == ModelStatus.Ready


@responses.activate
def test_save_model_no_refset(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': [],
        'status': ModelStatus.Pending
    }, status=201)

    with patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
               return_value=ModelStatus.Ready) as model_status_waiter_mock:
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    model_status_waiter_mock.assert_called_once_with([ModelStatus.Ready, ModelStatus.CreationFailed],
                                                     ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1
    assert reg_unsaved_batch_model.status == ModelStatus.Ready


@responses.activate
def test_save_model_provisioning_failed(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': [],
        'status': ModelStatus.Pending
    }, status=201)

    with pytest.raises(ArthurInternalError), \
            patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
                  return_value=ModelStatus.CreationFailed) as model_status_waiter_mock:
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_not_called()
    model_status_waiter_mock.assert_called_once_with([ModelStatus.Ready, ModelStatus.CreationFailed],
                                                     ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1
    assert reg_unsaved_batch_model.status == ModelStatus.CreationFailed

@responses.activate
def test_save_model_provisioning_failed_retry(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': [],
        'status': ModelStatus.Pending
    }, status=201)

    mock_post(f"/api/v3/models/{reg.TABULAR_MODEL_ID}/retry", response_body={}, status=200)

    with pytest.raises(ArthurInternalError), \
            patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
                  return_value=ModelStatus.CreationFailed):
        reg_unsaved_batch_model.save()

    with patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
               return_value=ModelStatus.Ready) as model_status_waiter_mock:
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    model_status_waiter_mock.assert_called_once_with([ModelStatus.Ready, ModelStatus.CreationFailed],
                                                     ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE)
    assert len(responses.calls) == 2
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1
    assert reg_unsaved_batch_model.status == ModelStatus.Ready

@responses.activate
def test_save_model_twice_error(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': [],
        'status': ModelStatus.Pending
    }, status=201)

    mock_post(f"/api/v3/models/{reg.TABULAR_MODEL_ID}/retry", response_body={}, status=200)

    with patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
               return_value=ModelStatus.Ready) as model_status_waiter_mock:
        reg_unsaved_batch_model.save()

    with pytest.raises(MethodNotApplicableError):
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    model_status_waiter_mock.assert_called_once_with([ModelStatus.Ready, ModelStatus.CreationFailed],
                                                     ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1
    assert reg_unsaved_batch_model.status == ModelStatus.Ready

@responses.activate
def test_save_model_twice_creating(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': [],
        'status': ModelStatus.Pending
    }, status=201)

    mock_post(f"/api/v3/models/{reg.TABULAR_MODEL_ID}/retry", response_body={}, status=200)

    with patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
               return_value=ModelStatus.Creating) as model_status_waiter_mock:
        reg_unsaved_batch_model.save()

    with pytest.raises(MethodNotApplicableError):
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    model_status_waiter_mock.assert_called_once_with([ModelStatus.Ready, ModelStatus.CreationFailed],
                                                     ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1
    assert reg_unsaved_batch_model.status == ModelStatus.Creating

@responses.activate
def test_save_model_twice_pending(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': [],
        'status': ModelStatus.Pending
    }, status=201)

    mock_post(f"/api/v3/models/{reg.TABULAR_MODEL_ID}/retry", response_body={}, status=200)

    with patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
               return_value=ModelStatus.Pending) as model_status_waiter_mock:
        reg_unsaved_batch_model.save()

    with pytest.raises(MethodNotApplicableError):
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    model_status_waiter_mock.assert_called_once_with([ModelStatus.Ready, ModelStatus.CreationFailed],
                                                     ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1
    assert reg_unsaved_batch_model.status == ModelStatus.Pending


@responses.activate
def test_save_model_provisioning_exception(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': [],
        'status': ModelStatus.Pending
    }, status=201)

    with pytest.raises(ArthurInternalError), \
            patch('arthurai.core.model_status_waiter.ModelStatusWaiter.wait_for_valid_model_status',
                  side_effect=ArthurInternalError) as model_status_waiter_mock:
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_not_called()
    model_status_waiter_mock.assert_called_once_with([ModelStatus.Ready, ModelStatus.CreationFailed],
                                                     ONBOARDING_SPINNER_MESSAGE, ONBOARDING_UPDATE_MESSAGE)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1
@responses.activate
def test_save_model_no_status_backwards_compatibility_check(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': []
    }, status=201)

    reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1

@responses.activate
def test_save_model_twice_no_status_backwards_compatibility_check(reg_unsaved_batch_model):
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'model_group_id': reg.TABULAR_MODEL_GROUP_ID,
        'version_sequence_num': 1,
        'attributes': []
    }, status=201)

    reg_unsaved_batch_model.save()

    with pytest.raises(MethodNotApplicableError):
        reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id == reg.TABULAR_MODEL_GROUP_ID
    assert reg_unsaved_batch_model.version_sequence_num == 1


@responses.activate
def test_save_model_binary_implicit_negative(initial_biclass_model):
    pred_gt_map = {mclass.DOG_PRED: mclass.DOG_GROUND_TRUTH, None: mclass.CAT_GROUND_TRUTH}
    model = initial_biclass_model
    model.build(data=SAMPLE_DATAFRAME_NO_NEG_PRED, pred_to_ground_truth_map=pred_gt_map,
                non_input_columns=[mclass.FLOAT_NONINPUT], set_reference_data=False)
    mock_post("/api/v3/models", response_body=MODEL_DATA_WITH_ATTRIBUTES_IMPLICIT_NEG, status=201)
    model.save()

    assert model.id == mclass.MODEL_ID
    assert_attributes_equal(mclass.MODEL_ATTRIBUTES_IMPLICIT_NEG, model.attributes)


@responses.activate
def test_save_model_without_model_group_id(reg_unsaved_batch_model):
    # Ensure models can still be saved even when the POST `models/` endpoint does not response with model_group_id
    # For testing backwards compatibility of model versioning SDK changes with pre-model-versioning API versions
    reg_unsaved_batch_model.reference_dataframe = reg.SAMPLE_DATAFRAME
    mock_set_ref_data = MagicMock()
    reg_unsaved_batch_model.set_reference_data = mock_set_ref_data
    mock_post("/api/v3/models", response_body={
        'id': reg.TABULAR_MODEL_ID,
        'attributes': []
    }, status=201)

    reg_unsaved_batch_model.save()

    mock_set_ref_data.assert_called_once_with(data=reg.SAMPLE_DATAFRAME)
    assert len(responses.calls) == 1
    assert os.getenv("ARTHUR_LAST_MODEL_ID") == reg.TABULAR_MODEL_ID
    assert reg_unsaved_batch_model.model_group_id is None
    assert reg_unsaved_batch_model.version_sequence_num is None


def test_add_attribute_positions(reg_unsaved_batch_model):
    reg_unsaved_batch_model.infer_schema(reg.SAMPLE_DATAFRAME[reg.INT_INPUT], stage=Stage.ModelPipelineInput)
    reg_unsaved_batch_model.infer_schema(reg.SAMPLE_DATAFRAME[reg.FLOAT_NONINPUT], stage=Stage.NonInputData)
    reg_unsaved_batch_model.add_regression_output_attributes({reg.PRED: reg.GROUND_TRUTH}, value_type=ValueType.Integer)
    # there is one actual attribute in each stage so should be two in each with duplicates
    expected_attr_positions = {0, 1}
    for stage in (Stage.ModelPipelineInput, Stage.NonInputData, Stage.PredictedValue, Stage.GroundTruth):
        actual_attr_positions = [attr.position for attr in reg_unsaved_batch_model.attributes if attr.stage == stage]
        assert len(actual_attr_positions) == len(expected_attr_positions)
        assert set(actual_attr_positions) == expected_attr_positions


@responses.activate
def test_find_hotspots(streaming_model):
    model = streaming_model

    response_body = {'data': [{'rules': {},
                               'gt_to_info': {'-1': {'class_f1': 0.8870523415977962,
                                                     'class_precision': 0.8298969072164949,
                                                     'class_recall': 0.9526627218934911,
                                                     'count': 169,
                                                     'pred_to_count': {'-1': 161, '1': 8}},
                                              '1': {'class_f1': 0.9356357927786499,
                                                    'class_precision': 0.9738562091503268,
                                                    'class_recall': 0.9003021148036254,
                                                    'count': 331,
                                                    'pred_to_count': {'-1': 33, '1': 298}}},
                               'precision': 0.9251979650966916,
                               'recall': 0.918,
                               'f1': 0.9192145862795214,
                               'accuracy': 0.918,
                               'impurity': 0.15055200000000002,
                               'n_samples': 500,
                               'feature': 'x2',
                               'cutoff': -3.560836390126483}]}

    mock_get(f"/api/v3/models/{model.id}/enrichments/hotspots/find?metric=f1&threshold=0.999&date=2021-08-04",
             response_body=response_body, status=200)
    model.find_hotspots(metric="f1", threshold=0.999, date="2021-08-04")
    assert len(responses.calls) == 1


@responses.activate
def test_query(streaming_model):
    query_body = {"property": "my_base_prop"}
    response_body = {'query_result': "base query response"}
    mock_post(f"/api/v3/models/{streaming_model.id}/inferences/query",
              response_body=response_body, status=200)

    streaming_model.query(query_body)

    assert len(responses.calls) == 1
    assert responses.calls[0].request.body == json.dumps(query_body)


@responses.activate
def test_query_psi_buckets(streaming_model):
    query_body = {"property": "my_psi_prop"}
    response_body = {'query_result': "psi bucket response"}
    mock_post(f"/api/v3/models/{streaming_model.id}/inferences/query/data_drift_psi_bucket_calculation_table",
              response_body=response_body, status=200)

    streaming_model.query(query_body, query_type="drift_psi_bucket_table")

    # assert request body was query
    assert len(responses.calls) == 1
    assert responses.calls[0].request.body == json.dumps(query_body)


@responses.activate
def test_query_bad_type(streaming_model):
    query_body = {"property": "my_prop"}
    with pytest.raises(UserValueError):
        streaming_model.query(query_body, query_type="foobar")


def test_model_group_from_dict():
    model_group = ArthurModelGroup.from_dict(MODEL_GROUP_JSON)

    assert model_group.id == MODEL_GROUP_ID
    assert model_group.name == MODEL_GROUP_NAME
    assert model_group.description == MODEL_GROUP_DESCRIPTION
    assert model_group.archived is MODEL_GROUP_ARCHIVED  # boolean value
    assert model_group.created_at == MODEL_GROUP_CREATED_AT
    assert model_group.updated_at == MODEL_GROUP_UPDATED_AT


def test_model_group_add_version():
    model = ArthurModel.from_dict(MODEL_RESPONSE_JSON[0])
    model_group = ArthurModelGroup.from_dict(MODEL_GROUP_JSON)
    label = "test_label"
    model_group.add_version(model, label)

    assert model.model_group_id == MODEL_GROUP_ID
    assert model.version_label == label


@responses.activate
def test_model_group_update_name(client):
    model_group = ArthurModelGroup.from_dict(MODEL_GROUP_JSON)
    model_group._update_client(client.client)
    model_group_new_name = "New Name"

    patch_response = MODEL_GROUP_JSON
    patch_response['name'] = model_group_new_name

    mock_patch(f"/api/v3/model_groups/{MODEL_GROUP_JSON['id']}",
               status=HTTPStatus.OK,
               response_body=patch_response)
    model_group.name = model_group_new_name

    assert model_group.name == model_group_new_name


@responses.activate
def test_model_group_update_description(client):
    model_group = ArthurModelGroup.from_dict(MODEL_GROUP_JSON)
    model_group._update_client(client.client)
    model_group_new_description = "New Description"

    patch_response = MODEL_GROUP_JSON
    patch_response['description'] = model_group_new_description


    mock_patch(f"/api/v3/model_groups/{MODEL_GROUP_JSON['id']}",
               status=HTTPStatus.OK,
               response_body=patch_response)
    model_group.description = model_group_new_description

    assert model_group.description == model_group_new_description


def test_get_attribute_wrong_stage():
    model_one = ArthurModel.from_json(model_response_json_strings[1])
    print(model_one.get_attributes(stage=Stage.ModelPipelineInput))
    # cannot call valid attribute name but invalid stage
    with pytest.raises(UserValueError):
        # correct stage is Stage.ModelPipelineInput
        model_one.get_attribute(name='attr_1', stage=Stage.GroundTruth)


def test_get_attribute_names_stage_none():
    model_one = ArthurModel.from_json(model_response_json_strings[1])
    attr_one = model_one.get_attribute_names(stage=None)
    assert len(attr_one) == 2

    model_two = ArthurModel.from_json(model_response_json_strings[2])
    attr_two = model_two.get_attribute_names(stage=None)
    assert len(attr_two) == 1
