"""Tests for the multi-modal image encoder (``palimseste.vision``).

Verifies:
  - encode_image returns a valid HV of the right dimensionality
  - similar images produce similar HVs (correlated)
  - dissimilar images produce quasi-orthogonal HVs
  - ImageMemory stores and retrieves by similarity
  - query_by_image finds the most similar stored image
"""

from __future__ import annotations

import numpy as np
import pytest

from palimseste import hv
from palimseste.vision import ImageEncoder, ImageMemory, encode_image


def _solid_image(color: tuple[int, int, int], size: int = 64) -> np.ndarray:
    """A solid-color image."""
    img = np.zeros((size, size, 3), dtype=np.uint8)
    img[:] = color
    return img


def _random_image(seed: int, size: int = 64) -> np.ndarray:
    rng = np.random.default_rng(seed)
    return rng.integers(0, 256, size=(size, size, 3), dtype=np.uint8)


class TestImageEncoder:
    def test_encode_returns_valid_hv(self):
        enc = ImageEncoder(D=2000, grid=8)
        img = _random_image(0)
        h = enc.encode(img)
        assert isinstance(h, hv.HV)
        assert h.D == 2000

    def test_similar_images_are_similar(self):
        enc = ImageEncoder(D=5000, grid=8)
        img1 = _solid_image((100, 50, 200))
        img2 = _solid_image((105, 55, 195))  # very close colors
        h1 = enc.encode(img1)
        h2 = enc.encode(img2)
        sim = hv.similarity(h1, h2)
        assert sim > 0.3, f"similar images should be correlated, got sim={sim}"

    def test_dissimilar_images_are_quasi_orthogonal(self):
        enc = ImageEncoder(D=5000, grid=8)
        img1 = _solid_image((0, 0, 0))    # black
        img2 = _solid_image((255, 255, 255))  # white
        # black vs white are opposite but share structure (solid), so not
        # fully orthogonal. Use random images instead.
        img3 = _random_image(1)
        img4 = _random_image(2)
        h3 = enc.encode(img3)
        h4 = enc.encode(img4)
        sim = hv.similarity(h3, h4)
        # random images with 8x8 grid share mean-RGB distribution structure,
        # so they're not fully orthogonal but should be well below 0.5
        assert sim < 0.5, f"random images should be weakly correlated, got sim={sim}"

    def test_same_image_same_hv(self):
        enc = ImageEncoder(D=3000, grid=8)
        img = _random_image(42)
        h1 = enc.encode(img)
        h2 = enc.encode(img)
        assert h1 == h2

    def test_encode_image_convenience(self):
        img = _random_image(0, size=32)
        h = encode_image(img, D=2000, grid=4)
        assert h.D == 2000

    def test_invalid_shape_raises(self):
        enc = ImageEncoder(D=1000)
        with pytest.raises(ValueError):
            enc.encode(np.zeros((10, 10), dtype=np.uint8))

    def test_small_image_works(self):
        enc = ImageEncoder(D=2000, grid=8)
        img = _random_image(0, size=4)  # smaller than grid
        h = enc.encode(img)
        assert h.D == 2000


class TestImageMemory:
    def test_store_and_retrieve(self):
        mem = ImageMemory(D=3000)
        img1 = _solid_image((255, 0, 0), size=32)
        img2 = _solid_image((0, 255, 0), size=32)
        mem.store(img1, label="red")
        mem.store(img2, label="green")
        results = mem.query_by_image(img1, top_k=1)
        assert len(results) == 1
        idx, sim, label = results[0]
        assert label == "red"
        assert sim > 0.5

    def test_retrieve_top_k(self):
        mem = ImageMemory(D=3000)
        for i in range(10):
            img = _solid_image((i * 25, 0, 0), size=32)
            mem.store(img, label=f"color_{i}")
        query = _solid_image((50, 0, 0), size=32)
        results = mem.query_by_image(query, top_k=3)
        assert len(results) == 3
        # results should be sorted by descending similarity
        assert results[0][1] >= results[1][1] >= results[2][1]

    def test_empty_memory(self):
        mem = ImageMemory(D=1000)
        assert mem.retrieve(hv.random_hv(1000)) == []

    def test_size(self):
        mem = ImageMemory(D=1000)
        assert mem.size == 0
        mem.store(_random_image(0, size=16))
        assert mem.size == 1