from textwrap import dedent
from unittest.mock import MagicMock
from uuid import uuid4

import pytest

from .datasource import Datasource
from .embedding_model import PretrainedEmbeddingModel
from .memoryset import ScoredMemoryset
from .regression_model import RegressionMetrics, RegressionModel
from .telemetry import RegressionPrediction


def test_create_model(regression_model: RegressionModel, scored_memoryset: ScoredMemoryset):
    assert regression_model is not None
    assert regression_model.name == "test_regression_model"
    assert regression_model.memoryset == scored_memoryset
    assert regression_model.memory_lookup_count == 3


def test_create_model_already_exists_error(scored_memoryset, regression_model: RegressionModel):
    with pytest.raises(ValueError):
        RegressionModel.create("test_regression_model", scored_memoryset)
    with pytest.raises(ValueError):
        RegressionModel.create("test_regression_model", scored_memoryset, if_exists="error")


def test_create_model_already_exists_return(scored_memoryset, regression_model: RegressionModel):
    with pytest.raises(ValueError):
        RegressionModel.create("test_regression_model", scored_memoryset, if_exists="open", memory_lookup_count=37)

    new_model = RegressionModel.create("test_regression_model", scored_memoryset, if_exists="open")
    assert new_model is not None
    assert new_model.name == "test_regression_model"
    assert new_model.memoryset == scored_memoryset
    assert new_model.memory_lookup_count == 3


def test_create_model_unauthenticated(unauthenticated_client, scored_memoryset: ScoredMemoryset):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            RegressionModel.create("test_regression_model", scored_memoryset)


def test_get_model(regression_model: RegressionModel):
    fetched_model = RegressionModel.open(regression_model.name)
    assert fetched_model is not None
    assert fetched_model.id == regression_model.id
    assert fetched_model.name == regression_model.name
    assert fetched_model.memory_lookup_count == 3
    assert fetched_model == regression_model


def test_get_model_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            RegressionModel.open("test_regression_model")


def test_get_model_invalid_input():
    with pytest.raises(ValueError, match="Invalid input"):
        RegressionModel.open("not valid id")


def test_get_model_not_found():
    with pytest.raises(LookupError):
        RegressionModel.open(str(uuid4()))


def test_get_model_unauthorized(unauthorized_client, regression_model: RegressionModel):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            RegressionModel.open(regression_model.name)


def test_list_models(regression_model: RegressionModel):
    models = RegressionModel.all()
    assert len(models) > 0
    assert any(model.name == regression_model.name for model in models)


def test_list_models_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            RegressionModel.all()


def test_list_models_unauthorized(unauthorized_client, regression_model: RegressionModel):
    with unauthorized_client.use():
        assert RegressionModel.all() == []


def test_memoryset_regression_models_property(regression_model: RegressionModel, scored_memoryset: ScoredMemoryset):
    models = scored_memoryset.regression_models
    assert len(models) > 0
    assert any(model.id == regression_model.id for model in models)


def test_update_model_attributes(regression_model: RegressionModel):
    regression_model.description = "New description"
    assert regression_model.description == "New description"

    regression_model.set(description=None)
    assert regression_model.description is None

    regression_model.set(locked=True)
    assert regression_model.locked is True

    regression_model.set(locked=False)
    assert regression_model.locked is False

    regression_model.lock()
    assert regression_model.locked is True

    regression_model.unlock()
    assert regression_model.locked is False


def test_delete_model(scored_memoryset: ScoredMemoryset):
    RegressionModel.create("regression_model_to_delete", ScoredMemoryset.open(scored_memoryset.name))
    assert RegressionModel.open("regression_model_to_delete")
    RegressionModel.drop("regression_model_to_delete")
    with pytest.raises(LookupError):
        RegressionModel.open("regression_model_to_delete")


def test_delete_model_unauthenticated(unauthenticated_client, regression_model: RegressionModel):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            RegressionModel.drop(regression_model.name)


def test_delete_model_not_found():
    with pytest.raises(LookupError):
        RegressionModel.drop(str(uuid4()))
    # ignores error if specified
    RegressionModel.drop(str(uuid4()), if_not_exists="ignore")


def test_delete_model_unauthorized(unauthorized_client, regression_model: RegressionModel):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            RegressionModel.drop(regression_model.name)


def test_delete_memoryset_before_model_constraint_violation(datasource: Datasource):
    memoryset = ScoredMemoryset.from_datasource("test_memoryset_delete_before_regression_model", datasource=datasource)
    RegressionModel.create("test_regression_model_delete_before_memoryset", memoryset)
    with pytest.raises(RuntimeError):
        ScoredMemoryset.drop(memoryset.id)


def test_delete_memoryset_with_model_cascade(datasource: Datasource):
    """Test that cascade=False prevents deletion and cascade=True allows it."""
    memoryset = ScoredMemoryset.from_datasource("test_memoryset_cascade_delete_regression", datasource=datasource)
    model = RegressionModel.create("test_regression_model_cascade_delete", memoryset)

    # Verify model exists
    assert RegressionModel.open(model.name) is not None

    # Without cascade, deletion should fail
    with pytest.raises(RuntimeError):
        ScoredMemoryset.drop(memoryset.id, cascade=False)

    # Model should still exist
    assert RegressionModel.exists(model.name)

    # With cascade, deletion should succeed
    ScoredMemoryset.drop(memoryset.id, cascade=True)

    # Model should be deleted along with the memoryset
    assert not RegressionModel.exists(model.name)
    assert not ScoredMemoryset.exists(memoryset.name)


def test_evaluate_datasource(regression_model: RegressionModel, eval_datasource: Datasource):
    """Test that model evaluation with a datasource works."""
    result = regression_model.evaluate(eval_datasource)

    assert isinstance(result, RegressionMetrics)
    assert result.mae == pytest.approx(0.4)
    assert 0.0 <= result.mse <= 1.0
    assert 0.0 <= result.rmse <= 1.0
    assert result.r2 is not None

    assert isinstance(result.anomaly_score_mean, float)
    assert isinstance(result.anomaly_score_median, float)
    assert isinstance(result.anomaly_score_variance, float)
    assert -1.0 <= result.anomaly_score_mean <= 1.0
    assert -1.0 <= result.anomaly_score_median <= 1.0
    assert -1.0 <= result.anomaly_score_variance <= 1.0


def test_evaluate_local_dataset(regression_model: RegressionModel, eval_data: list[dict]):
    datasets = pytest.importorskip("datasets")
    eval_dataset = datasets.Dataset.from_list(eval_data)
    result = regression_model.evaluate(eval_dataset)

    assert isinstance(result, RegressionMetrics)
    assert isinstance(result.mae, float)
    assert isinstance(result.mse, float)
    assert isinstance(result.rmse, float)
    assert result.r2 is not None


def test_evaluate_local_list(regression_model: RegressionModel, eval_data: list[dict]):
    result = regression_model.evaluate(eval_data)
    assert result is not None
    assert isinstance(result, RegressionMetrics)
    assert isinstance(result.mae, float)
    assert isinstance(result.mse, float)
    assert isinstance(result.rmse, float)


def test_evaluate_datasource_with_nones_raises_error(regression_model: RegressionModel, datasource: Datasource):
    with pytest.raises(ValueError):
        regression_model.evaluate(datasource, record_predictions=True, tags={"test"})


def test_evaluate_with_telemetry(regression_model, eval_data: list[dict], eval_datasource: Datasource):
    result = regression_model.evaluate(eval_datasource, record_predictions=True, tags={"test"}, batch_size=2)
    assert result is not None
    assert isinstance(result, RegressionMetrics)
    predictions = regression_model.predictions(tag="test", batch_size=100, sort=[("timestamp", "asc")])
    assert len(predictions) == 4
    assert all(p.tags == {"test"} for p in predictions)
    assert all(p.expected_score is not None for p in predictions)
    prediction_expected_scores = [p.expected_score for p in predictions]
    eval_expected_scores = [d["score"] for d in eval_data]
    assert prediction_expected_scores == pytest.approx(eval_expected_scores, abs=1e-3)


def test_evaluate_with_partition_column_include_global(partitioned_regression_model: RegressionModel):
    result = partitioned_regression_model.evaluate(
        [
            {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
            {"value": "cats are cute", "score": 0.9, "partition_id": "p1"},
        ],
        partition_column="partition_id",
        partition_filter_mode="include_global",
    )
    assert result is not None
    assert isinstance(result, RegressionMetrics)


def test_evaluate_with_partition_column_exclude_global(partitioned_regression_model: RegressionModel):
    result = partitioned_regression_model.evaluate(
        [
            {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
            {"value": "cats are cute", "score": 0.9, "partition_id": "p1"},
        ],
        partition_column="partition_id",
        partition_filter_mode="exclude_global",
    )
    assert result is not None
    assert isinstance(result, RegressionMetrics)


def test_evaluate_with_partition_column_only_global(partitioned_regression_model: RegressionModel):
    result = partitioned_regression_model.evaluate(
        [
            {"value": "cats are independent animals", "score": 0.9, "partition_id": None},
            {"value": "i love the beach", "score": 0.5, "partition_id": None},
        ],
        partition_column="partition_id",
        partition_filter_mode="only_global",
    )
    assert result is not None
    assert isinstance(result, RegressionMetrics)


def test_evaluate_with_partition_column_ignore_partitions(partitioned_regression_model: RegressionModel):
    result = partitioned_regression_model.evaluate(
        [
            {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
            {"value": "cats are cute", "score": 0.9, "partition_id": "p2"},
        ],
        partition_column="partition_id",
        partition_filter_mode="ignore_partitions",
    )
    assert result is not None
    assert isinstance(result, RegressionMetrics)


def test_evaluate_with_partition_column_datasource(partitioned_regression_model: RegressionModel):
    eval_datasource = Datasource.from_list(
        "eval_partition_datasource_reg",
        [
            {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
            {"value": "cats are cute", "score": 0.9, "partition_id": "p1"},
            {"value": "homemade soup recipes", "score": 0.1, "partition_id": "p2"},
            {"value": "cats purr when happy", "score": 0.9, "partition_id": "p2"},
        ],
    )
    result = partitioned_regression_model.evaluate(
        eval_datasource,
        partition_column="partition_id",
        partition_filter_mode="exclude_global",
    )
    assert result is not None
    assert isinstance(result, RegressionMetrics)


def test_predict(regression_model: RegressionModel):
    predictions = regression_model.predict(["Do you love soup?", "Are cats cute?"], batch_size=1)
    assert len(predictions) == 2
    assert predictions[0].prediction_id is not None
    assert predictions[1].prediction_id is not None
    assert predictions[0].score is not None
    assert predictions[1].score is not None
    assert predictions[0].score == pytest.approx(0.1)
    assert predictions[1].score == pytest.approx(0.9)
    assert 0 <= predictions[0].confidence <= 1
    assert 0 <= predictions[1].confidence <= 1


def test_regression_prediction_has_no_score(regression_model: RegressionModel):
    """Ensure optional score is None for regression predictions."""
    prediction = regression_model.predict("This beach is amazing!")
    assert isinstance(prediction, RegressionPrediction)
    assert prediction.score is None


def test_predict_unauthenticated(unauthenticated_client, regression_model: RegressionModel):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            regression_model.predict(["This is excellent!", "This is terrible!"])


def test_predict_unauthorized(unauthorized_client, regression_model: RegressionModel):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            regression_model.predict(["This is excellent!", "This is terrible!"])


def test_predict_constraint_violation(scored_memoryset: ScoredMemoryset):
    model = RegressionModel.create(
        "test_regression_model_lookup_count_too_high",
        scored_memoryset,
        memory_lookup_count=scored_memoryset.length + 2,
    )
    with pytest.raises(RuntimeError):
        model.predict("test")


def test_predict_with_prompt(regression_model: RegressionModel):
    """Test that prompt parameter is properly passed through to predictions"""
    # Test with an instruction-supporting embedding model if available
    prediction_with_prompt = regression_model.predict(
        "This product is amazing!", prompt="Represent this text for rating prediction:"
    )
    prediction_without_prompt = regression_model.predict("This product is amazing!")

    # Both should work and return valid predictions
    assert prediction_with_prompt.score is not None
    assert prediction_without_prompt.score is not None
    assert 0 <= prediction_with_prompt.confidence <= 1
    assert 0 <= prediction_without_prompt.confidence <= 1


def test_predict_with_partition_id(partitioned_regression_model: RegressionModel):
    """Test predict with a specific partition_id"""
    # Predict with partition_id p1 - should use memories from p1
    prediction = partitioned_regression_model.predict("soup", partition_id="p1", partition_filter_mode="exclude_global")
    assert prediction.score is not None
    assert 0 <= prediction.confidence <= 1

    # Predict with partition_id p2 - should use memories from p2
    prediction_p2 = partitioned_regression_model.predict(
        "cats", partition_id="p2", partition_filter_mode="exclude_global"
    )
    assert prediction_p2.score is not None
    assert 0 <= prediction_p2.confidence <= 1


def test_predict_with_partition_id_include_global(partitioned_regression_model: RegressionModel):
    """Test predict with partition_id and include_global mode (default)"""
    # Predict with partition_id p1 and include_global (default) - should include both p1 and global memories
    prediction = partitioned_regression_model.predict("soup", partition_id="p1", partition_filter_mode="include_global")
    assert prediction.score is not None
    assert 0 <= prediction.confidence <= 1


def test_predict_with_partition_id_exclude_global(partitioned_regression_model: RegressionModel):
    """Test predict with partition_id and exclude_global mode"""
    # Predict with partition_id p1 and exclude_global - should only use p1 memories
    prediction = partitioned_regression_model.predict("soup", partition_id="p1", partition_filter_mode="exclude_global")
    assert prediction.score is not None
    assert 0 <= prediction.confidence <= 1


def test_predict_with_partition_id_only_global(partitioned_regression_model: RegressionModel):
    """Test predict with partition_filter_mode only_global"""
    # Predict with only_global mode - should only use global memories
    prediction = partitioned_regression_model.predict("cats", partition_filter_mode="only_global")
    assert prediction.score is not None
    assert 0 <= prediction.confidence <= 1


def test_predict_with_partition_id_ignore_partitions(partitioned_regression_model: RegressionModel):
    """Test predict with partition_filter_mode ignore_partitions"""
    # Predict with ignore_partitions mode - should ignore partition filtering
    prediction = partitioned_regression_model.predict("soup", partition_filter_mode="ignore_partitions")
    assert prediction.score is not None
    assert 0 <= prediction.confidence <= 1


def test_predict_batch_with_partition_id(partitioned_regression_model: RegressionModel):
    """Test batch predict with partition_id"""
    # Batch predict with partition_id p1
    predictions = partitioned_regression_model.predict(
        ["soup is good", "cats are cute"],
        partition_id="p1",
        partition_filter_mode="exclude_global",
    )
    assert len(predictions) == 2
    assert all(p.score is not None for p in predictions)
    assert all(0 <= p.confidence <= 1 for p in predictions)


def test_predict_batch_with_list_of_partition_ids(partitioned_regression_model: RegressionModel):
    """Test batch predict with a list of partition_ids (one for each query input)"""
    # Batch predict with a list of partition_ids - one for each input
    # First input uses p1, second input uses p2
    predictions = partitioned_regression_model.predict(
        ["soup is good", "cats are cute"],
        partition_id=["p1", "p2"],
        partition_filter_mode="exclude_global",
    )
    assert len(predictions) == 2
    assert all(p.score is not None for p in predictions)
    assert all(0 <= p.confidence <= 1 for p in predictions)

    # Verify that predictions were made using the correct partitions
    # Each prediction should use memories from its respective partition
    assert predictions[0].input_value == "soup is good"
    assert predictions[1].input_value == "cats are cute"


def test_record_prediction_feedback(regression_model: RegressionModel):
    predictions = regression_model.predict(["This is excellent!", "This is terrible!"])
    expected_scores = [0.9, 0.1]
    regression_model.record_feedback(
        {
            "prediction_id": p.prediction_id,
            "category": "accurate",
            "value": abs(p.score - expected_score) < 0.2 if p.score is not None else False,
        }
        for expected_score, p in zip(expected_scores, predictions)
    )


def test_record_prediction_feedback_missing_category(regression_model: RegressionModel):
    prediction = regression_model.predict("This is excellent!")
    with pytest.raises(ValueError):
        regression_model.record_feedback({"prediction_id": prediction.prediction_id, "value": True})


def test_record_prediction_feedback_invalid_value(regression_model: RegressionModel):
    prediction = regression_model.predict("This is excellent!")
    with pytest.raises(ValueError, match=r"Invalid input.*"):
        regression_model.record_feedback(
            {"prediction_id": prediction.prediction_id, "category": "accurate", "value": "invalid"}
        )


def test_record_prediction_feedback_invalid_prediction_id(regression_model: RegressionModel):
    with pytest.raises(ValueError, match=r"Invalid input.*"):
        regression_model.record_feedback({"prediction_id": "invalid", "category": "accurate", "value": True})


def test_predict_with_memoryset_override(regression_model: RegressionModel, data: list[dict]):
    # Create a memoryset with different scores (inverted)
    inverted_scored_memoryset = ScoredMemoryset.from_list(
        "test_memoryset_inverted_scores",
        [{"value": x["value"], "score": (2.0 - x["score"]) if x["score"] is not None else None} for x in data],
        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
    )
    original_predictions = regression_model.predict(["This is excellent!", "This is terrible!"])

    with regression_model.use_memoryset(inverted_scored_memoryset):
        override_predictions = regression_model.predict(
            ["This is excellent!", "This is terrible!"], tags={"override_test"}
        )
        # With inverted scores, the predictions should be different
        assert original_predictions[0].score is not None
        assert original_predictions[1].score is not None
        assert override_predictions[0].score is not None
        assert override_predictions[1].score is not None
        assert abs(override_predictions[0].score - original_predictions[0].score) > 0.1
        assert abs(override_predictions[1].score - original_predictions[1].score) > 0.1
        assert override_predictions[0].memoryset == inverted_scored_memoryset
        assert override_predictions[1].memoryset == inverted_scored_memoryset

    # After exiting context, predictions should be back to normal
    new_predictions = regression_model.predict(["This is excellent!", "This is terrible!"], tags={"no_override_test"})
    assert new_predictions[0].score is not None
    assert new_predictions[1].score is not None
    assert original_predictions[0].score is not None
    assert original_predictions[1].score is not None
    assert abs(new_predictions[0].score - original_predictions[0].score) < 0.1
    assert abs(new_predictions[1].score - original_predictions[1].score) < 0.1
    assert new_predictions[0].memoryset == regression_model.memoryset
    assert new_predictions[1].memoryset == regression_model.memoryset

    # Verify historical predictions have the correct memoryset from telemetry
    fetched_override_predictions = regression_model.predictions(tag="override_test")
    assert len(fetched_override_predictions) == 2
    assert fetched_override_predictions[0].memoryset == inverted_scored_memoryset
    assert fetched_override_predictions[1].memoryset == inverted_scored_memoryset

    fetched_no_override_predictions = regression_model.predictions(tag="no_override_test")
    assert len(fetched_no_override_predictions) == 2
    assert fetched_no_override_predictions[0].memoryset == regression_model.memoryset
    assert fetched_no_override_predictions[1].memoryset == regression_model.memoryset


def test_predict_with_expected_scores(regression_model: RegressionModel):
    prediction = regression_model.predict("This is excellent!", expected_scores=0.9)
    assert prediction.expected_score == 0.9


def test_regression_prediction_update(regression_model: RegressionModel):
    prediction = regression_model.predict("Test input", expected_scores=3.5)
    assert prediction.expected_score == 3.5
    assert prediction.tags == set()

    # Update expected score
    prediction.update(expected_score=4.5)
    assert prediction.expected_score == 4.5

    # Add tags
    prediction.update(tags={"test", "updated"})
    assert prediction.tags == {"test", "updated"}

    # Clear both
    prediction.update(expected_score=None, tags=None)
    assert prediction.expected_score is None
    assert prediction.tags == set()


def test_last_prediction_with_batch(regression_model: RegressionModel):
    predictions = regression_model.predict(["This is excellent!", "This is terrible!"])
    assert regression_model.last_prediction is not None
    assert regression_model.last_prediction.prediction_id == predictions[-1].prediction_id
    assert regression_model.last_prediction.input_value == "This is terrible!"
    assert regression_model._last_prediction_was_batch is True


def test_last_prediction_with_single(regression_model: RegressionModel):
    # Test that last_prediction is updated correctly with single prediction
    prediction = regression_model.predict("This is excellent!")
    assert regression_model.last_prediction is not None
    assert regression_model.last_prediction.prediction_id == prediction.prediction_id
    assert regression_model.last_prediction.input_value == "This is excellent!"
    assert regression_model._last_prediction_was_batch is False


def test_batch_predict(regression_model: RegressionModel):
    """Test batch predictions"""
    predictions = regression_model.predict(["test input 1", "test input 2", "test input 3"])
    assert len(predictions) == 3
    assert all(isinstance(pred, RegressionPrediction) for pred in predictions)


def test_batch_predict_with_expected_scores(regression_model: RegressionModel):
    """Test batch predictions with expected scores"""
    predictions = regression_model.predict(["input 1", "input 2"], expected_scores=[0.5, 0.8])
    assert len(predictions) == 2
    assert all(isinstance(pred, RegressionPrediction) for pred in predictions)


def test_use_memoryset(regression_model: RegressionModel, scored_memoryset: ScoredMemoryset):
    # Test that predictions work with a memoryset
    predictions = regression_model.predict(["This is excellent!", "This is terrible!"])
    assert len(predictions) == 2
    assert all(isinstance(pred, RegressionPrediction) for pred in predictions)
    assert all(0 <= pred.confidence <= 1 for pred in predictions)

    # Test that predictions work with a different memoryset
    with regression_model.use_memoryset(scored_memoryset):
        predictions = regression_model.predict(["This is excellent!", "This is terrible!"])
        assert len(predictions) == 2
        assert all(isinstance(pred, RegressionPrediction) for pred in predictions)
        assert all(0 <= pred.confidence <= 1 for pred in predictions)


def test_drop(regression_model):
    """Test that model drop works."""
    name = regression_model.name
    RegressionModel.drop(name)
    assert not RegressionModel.exists(name)


def test_predict_with_empty_partition(fully_partitioned_regression_resources):
    datasource, memoryset, regression_model = fully_partitioned_regression_resources

    assert memoryset.length == 15

    with pytest.raises(RuntimeError, match="lookup failed to return the correct number of memories"):
        regression_model.predict("i love cats", partition_filter_mode="only_global")

    with pytest.raises(RuntimeError, match="lookup failed to return the correct number of memories"):
        regression_model.predict("i love cats", partition_filter_mode="exclude_global", partition_id="p_does_not_exist")

    with pytest.raises(RuntimeError, match="lookup failed to return the correct number of memories"):
        regression_model.evaluate(datasource, partition_filter_mode="only_global")


def test_regression_model_repr(regression_model: RegressionModel):
    assert repr(regression_model) == dedent(
        """\
        RegressionModel({
            name: 'test_regression_model',
            memory_lookup_count: 3,
            memoryset: ScoredMemoryset({
                name: 'test_scored_memoryset',
                length: 22,
                embedding_model: PretrainedEmbeddingModel({ name: GTE_BASE, embedding_dim: 768, max_seq_length: 8192, num_params: 137M }),
            }),
        })"""
    )


def test_regression_metrics_repr():
    metrics = RegressionMetrics(
        {
            "coverage": 1.0,
            "mse": 0.1600,
            "rmse": 0.4,
            "mae": 0.4,
            "r2": -0.3333,
            "explained_variance": 0.0,
            "loss": 0.16,
            "anomaly_score_mean": -0.1234,
            "anomaly_score_median": -0.1,
            "anomaly_score_variance": 0.0567,
        }
    )
    assert repr(metrics) == dedent(
        """\
        RegressionMetrics({
            mae: 0.4000,
            rmse: 0.4000,
            r2: -0.3333,
            anomaly_score: -0.1234 ± 0.0567,
        })"""
    )


def test_regression_metrics_repr_without_optional_fields():
    metrics = RegressionMetrics(
        {
            "coverage": 1.0,
            "mse": 0.16,
            "rmse": 0.4,
            "mae": 0.4,
            "r2": 0.75,
            "explained_variance": 0.5,
            "loss": 0.16,
        }
    )
    assert repr(metrics) == dedent(
        """\
        RegressionMetrics({
            mae: 0.4000,
            rmse: 0.4000,
            r2: 0.7500,
        })"""
    )


def test_regression_prediction_repr():
    prediction = RegressionPrediction(
        None,
        label=None,
        label_name=None,
        score=0.1567,
        confidence=0.9234,
        anomaly_score=-0.0567,
        memoryset=MagicMock(),
        model=MagicMock(),
        input_value="Do you love soup?",
    )
    assert repr(prediction) == (
        "RegressionPrediction({ score: 0.16, confidence: 0.92, anomaly_score: -0.06, input_value: 'Do you love soup?' })"
    )
