import pytest
from pydantic import ValidationError

from backend.app.models.campaign import HeatmapRequest
from backend.app.services.heatmap import build_heatmap


def _identity_covariance() -> list[list[float]]:
    return [[1.0, 0.0], [0.0, 1.0]]


def test_build_heatmap_ranks_by_distance_and_softmax_probability() -> None:
    result = build_heatmap(
        campaign_vector=[0.0, 0.0],
        centroids=[[0.0, 0.0], [3.0, 0.0]],
        covariances=[_identity_covariance(), _identity_covariance()],
        cluster_top_doctors={"0": ["Dr_A"], "1": ["Dr_B"]},
    )

    ranking = result["ranking"]
    assert len(ranking) == 2
    assert ranking[0]["cluster_id"] == 0
    assert ranking[0]["distance"] < ranking[1]["distance"]
    assert ranking[0]["probability"] > ranking[1]["probability"]
    assert abs(sum(item["probability"] for item in ranking) - 1.0) < 1e-9


def test_heatmap_request_rejects_invalid_dimensions() -> None:
    with pytest.raises(ValidationError):
        HeatmapRequest(
            campaign_vector=[1.0, 2.0],
            centroids=[[1.0, 2.0, 3.0]],
            covariances=[[[1.0, 0.0], [0.0, 1.0]]],
        )


def test_heatmap_request_accepts_valid_payload() -> None:
    payload = HeatmapRequest(
        campaign_vector=[0.0, 0.0],
        centroids=[[0.0, 0.0], [3.0, 0.0]],
        covariances=[_identity_covariance(), _identity_covariance()],
        cluster_top_doctors={"0": ["Dr_A"], "1": ["Dr_B"]},
    )

    result = build_heatmap(
        payload.campaign_vector,
        payload.centroids,
        payload.covariances,
        payload.cluster_top_doctors,
    )
    assert result["ranking"][0]["cluster_id"] == 0
    assert result["ranking"][0]["top_doctors"] == ["Dr_A"]