wled-screen-controller-mixed/server/tests/test_calibration.py

"""Tests for calibration system."""

import numpy as np
import pytest

from wled_controller.core.calibration import (
    CalibrationSegment,
    CalibrationConfig,
    PixelMapper,
    create_default_calibration,
    calibration_from_dict,
    calibration_to_dict,
)
from wled_controller.core.screen_capture import BorderPixels


def test_calibration_segment():
    """Test calibration segment creation."""
    segment = CalibrationSegment(
        edge="top",
        led_start=0,
        led_count=40,
        reverse=False,
    )

    assert segment.edge == "top"
    assert segment.led_start == 0
    assert segment.led_count == 40
    assert segment.reverse is False


def test_calibration_config_validation():
    """Test calibration configuration validation."""
    segments = [
        CalibrationSegment(edge="bottom", led_start=0, led_count=40),
        CalibrationSegment(edge="right", led_start=40, led_count=30),
        CalibrationSegment(edge="top", led_start=70, led_count=40),
        CalibrationSegment(edge="left", led_start=110, led_count=40),
    ]

    config = CalibrationConfig(
        layout="clockwise",
        start_position="bottom_left",
        segments=segments,
    )

    assert config.validate() is True
    assert config.get_total_leds() == 150


def test_calibration_config_duplicate_edges():
    """Test validation fails with duplicate edges."""
    segments = [
        CalibrationSegment(edge="top", led_start=0, led_count=40),
        CalibrationSegment(edge="top", led_start=40, led_count=40),  # Duplicate
    ]

    config = CalibrationConfig(
        layout="clockwise",
        start_position="bottom_left",
        segments=segments,
    )

    with pytest.raises(ValueError, match="Duplicate edges"):
        config.validate()


def test_calibration_config_overlapping_indices():
    """Test validation fails with overlapping LED indices."""
    segments = [
        CalibrationSegment(edge="bottom", led_start=0, led_count=50),
        CalibrationSegment(edge="right", led_start=40, led_count=30),  # Overlaps
    ]

    config = CalibrationConfig(
        layout="clockwise",
        start_position="bottom_left",
        segments=segments,
    )

    with pytest.raises(ValueError, match="overlap"):
        config.validate()


def test_calibration_config_invalid_led_count():
    """Test validation fails with invalid LED counts."""
    segments = [
        CalibrationSegment(edge="top", led_start=0, led_count=0),  # Invalid
    ]

    config = CalibrationConfig(
        layout="clockwise",
        start_position="bottom_left",
        segments=segments,
    )

    with pytest.raises(ValueError):
        config.validate()


def test_get_segment_for_edge():
    """Test getting segment by edge name."""
    segments = [
        CalibrationSegment(edge="bottom", led_start=0, led_count=40),
        CalibrationSegment(edge="right", led_start=40, led_count=30),
    ]

    config = CalibrationConfig(
        layout="clockwise",
        start_position="bottom_left",
        segments=segments,
    )

    bottom_seg = config.get_segment_for_edge("bottom")
    assert bottom_seg is not None
    assert bottom_seg.led_count == 40

    missing_seg = config.get_segment_for_edge("top")
    assert missing_seg is None


def test_pixel_mapper_initialization():
    """Test pixel mapper initialization."""
    config = create_default_calibration(150)
    mapper = PixelMapper(config, interpolation_mode="average")

    assert mapper.calibration == config
    assert mapper.interpolation_mode == "average"


def test_pixel_mapper_invalid_mode():
    """Test pixel mapper with invalid interpolation mode."""
    config = create_default_calibration(150)

    with pytest.raises(ValueError):
        PixelMapper(config, interpolation_mode="invalid")


def test_pixel_mapper_map_border_to_leds():
    """Test mapping border pixels to LED colors."""
    config = create_default_calibration(40)  # 10 per edge
    mapper = PixelMapper(config)

    # Create test border pixels (all red)
    border_pixels = BorderPixels(
        top=np.full((10, 100, 3), [255, 0, 0], dtype=np.uint8),
        right=np.full((100, 10, 3), [0, 255, 0], dtype=np.uint8),
        bottom=np.full((10, 100, 3), [0, 0, 255], dtype=np.uint8),
        left=np.full((100, 10, 3), [255, 255, 0], dtype=np.uint8),
    )

    led_colors = mapper.map_border_to_leds(border_pixels)

    assert len(led_colors) == 40
    assert all(isinstance(c, tuple) and len(c) == 3 for c in led_colors)

    # Verify colors are reasonable (allowing for some rounding)
    # Bottom LEDs should be mostly blue
    bottom_color = led_colors[0]
    assert bottom_color[2] > 200  # Blue channel high

    # Top LEDs should be mostly red
    top_segment = config.get_segment_for_edge("top")
    top_color = led_colors[top_segment.led_start]
    assert top_color[0] > 200  # Red channel high


def test_pixel_mapper_test_calibration():
    """Test calibration testing pattern."""
    config = create_default_calibration(100)
    mapper = PixelMapper(config)

    # Test top edge
    led_colors = mapper.test_calibration("top", (255, 0, 0))

    assert len(led_colors) == 100

    # Top edge should be lit (red)
    top_segment = config.get_segment_for_edge("top")
    top_leds = led_colors[top_segment.led_start:top_segment.led_start + top_segment.led_count]
    assert all(color == (255, 0, 0) for color in top_leds)

    # Other LEDs should be off
    other_leds = led_colors[:top_segment.led_start]
    assert all(color == (0, 0, 0) for color in other_leds)


def test_pixel_mapper_test_calibration_invalid_edge():
    """Test calibration testing with invalid edge."""
    config = CalibrationConfig(
        layout="clockwise",
        start_position="bottom_left",
        segments=[
            CalibrationSegment(edge="bottom", led_start=0, led_count=40),
        ],
    )
    mapper = PixelMapper(config)

    with pytest.raises(ValueError):
        mapper.test_calibration("top", (255, 0, 0))  # Top not in config


def test_create_default_calibration():
    """Test creating default calibration."""
    config = create_default_calibration(150)

    assert config.layout == "clockwise"
    assert config.start_position == "bottom_left"
    assert len(config.segments) == 4
    assert config.get_total_leds() == 150

    # Check all edges are present
    edges = {seg.edge for seg in config.segments}
    assert edges == {"top", "right", "bottom", "left"}


def test_create_default_calibration_small_count():
    """Test default calibration with small LED count."""
    config = create_default_calibration(4)
    assert config.get_total_leds() == 4


def test_create_default_calibration_invalid():
    """Test default calibration with invalid LED count."""
    with pytest.raises(ValueError):
        create_default_calibration(3)  # Too few LEDs


def test_calibration_from_dict():
    """Test creating calibration from dictionary."""
    data = {
        "layout": "clockwise",
        "start_position": "bottom_left",
        "segments": [
            {"edge": "bottom", "led_start": 0, "led_count": 40, "reverse": False},
            {"edge": "right", "led_start": 40, "led_count": 30, "reverse": False},
        ],
    }

    config = calibration_from_dict(data)

    assert config.layout == "clockwise"
    assert config.start_position == "bottom_left"
    assert len(config.segments) == 2
    assert config.get_total_leds() == 70


def test_calibration_from_dict_missing_field():
    """Test calibration from dict with missing field."""
    data = {
        "layout": "clockwise",
        # Missing start_position
        "segments": [],
    }

    with pytest.raises(ValueError):
        calibration_from_dict(data)


def test_calibration_to_dict():
    """Test converting calibration to dictionary."""
    config = create_default_calibration(100)
    data = calibration_to_dict(config)

    assert "layout" in data
    assert "start_position" in data
    assert "segments" in data
    assert isinstance(data["segments"], list)
    assert len(data["segments"]) == 4


def test_calibration_round_trip():
    """Test converting calibration to dict and back."""
    original = create_default_calibration(120)
    data = calibration_to_dict(original)
    restored = calibration_from_dict(data)

    assert restored.layout == original.layout
    assert restored.start_position == original.start_position
    assert len(restored.segments) == len(original.segments)
    assert restored.get_total_leds() == original.get_total_leds()