from __future__ import annotations

import numpy as np
import pytest

from idri import IdentifierEncoder, compose_vectors, cosine_similarity, l2_normalize


def test_compose_vectors_weighted_sum_and_normalization() -> None:
    enc = IdentifierEncoder(dim=256)
    v1 = enc.encode("omnisphere wm")
    v2 = enc.encode("omnisphere")

    composed = compose_vectors([(v1, 0.7), (v2, 0.3)], normalize=True)
    manual = l2_normalize((0.7 * v1 + 0.3 * v2).astype(np.float32))

    assert np.allclose(composed, manual, atol=1e-6)
    assert np.isclose(np.linalg.norm(composed), 1.0, atol=1e-6)


def test_compose_vectors_empty_input_returns_zero_vector() -> None:
    out = compose_vectors([], dim=64, normalize=True)
    assert out.shape == (64,)
    assert np.allclose(out, np.zeros(64, dtype=np.float32))


def test_compose_vectors_empty_input_without_dim_raises() -> None:
    with pytest.raises(ValueError):
        compose_vectors([], normalize=True)


def test_compose_vectors_dtype_output() -> None:
    enc = IdentifierEncoder(dim=64)
    v = enc.encode("abc")
    out = compose_vectors([(v, 1.0)], dtype="float64")
    assert out.dtype == np.float64


def test_vector_level_cosine_for_composed_vectors() -> None:
    enc = IdentifierEncoder(dim=256)
    va = compose_vectors([(enc.encode("drill set"), 0.5), (enc.encode("hammer"), 0.5)])
    vb = compose_vectors([(enc.encode("drill set"), 0.7), (enc.encode("hammer"), 0.3)])
    score = cosine_similarity(va, vb, assume_normalized=True)
    assert 0.0 <= score <= 1.0
