"""Unit tests for the simulated IDS camera frame source (_SimIDSBackend /
make_rotating_test_pattern / make_speckle_test_pattern /
SimIDSCamera's optional rotation coupling)."""

import numpy as np
import pytest

from csaxs_bec.devices.sim.sim_cameras import (
    SimIDSCamera,
    _SimIDSBackend,
    generate_speckle_grains,
    make_rotating_test_pattern,
    make_speckle_test_pattern,
    make_test_pattern,
)
from csaxs_bec.devices.sim.sim_galil import SimGalilState
from csaxs_bec.devices.sim.sim_socket import SimStateRegistry


@pytest.fixture(autouse=True)
def reset_sim_state():
    """Avoid cross-test leakage of shared (host, port)-keyed simulation state."""
    SimStateRegistry.reset()
    yield
    SimStateRegistry.reset()


def test_sim_ids_backend_default_matches_static_pattern():
    """No rotation_coupling (the default) -> byte-identical to today's static frame.

    Regression guard: this is what keeps flomni's cam_xeye and other existing
    SimIDSCamera users completely unaffected by the new opt-in feature.
    """
    backend = _SimIDSBackend(width=64, height=48, rgb=True)
    expected = make_test_pattern(48, 64, rgb=True)
    assert np.array_equal(backend.get_image_data(), expected)
    # calling it again returns the exact same cached frame, not a fresh one
    assert np.array_equal(backend.get_image_data(), expected)


@pytest.mark.parametrize("angle_deg", [30.0, 120.0, 210.0, 300.0])
def test_make_rotating_test_pattern_blob_position(angle_deg):
    height, width = 200, 300
    radius_px = 60.0
    frame = make_rotating_test_pattern(
        height, width, rgb=False, angle_deg=angle_deg, radius_px=radius_px
    )

    cy, cx = height / 2.0, width / 2.0
    angle_rad = np.deg2rad(angle_deg)
    expected_y = cy - radius_px * np.sin(angle_rad)
    expected_x = cx + radius_px * np.cos(angle_rad)

    # mask out the fixed center crosshair so it can't win the argmax
    masked = frame.astype(np.int32)
    masked[int(cy) - 2 : int(cy) + 3, :] = 0
    masked[:, int(cx) - 2 : int(cx) + 3] = 0
    py, px = np.unravel_index(np.argmax(masked), masked.shape)

    assert abs(py - expected_y) <= 2
    assert abs(px - expected_x) <= 2


def test_sim_ids_backend_rotation_coupling_tracks_live_angle():
    """_SimIDSBackend renders a speckle field (its actual default rotation-
    coupled pattern) tracking the live simulated angle, regenerated fresh
    per call rather than cached."""
    host, port, axis_letter = "test-galil.example", 9000, "C"
    galil = SimStateRegistry.get(SimGalilState, host, port)
    ax = galil.axis(ord(axis_letter.lower()) - 97)
    ax.stppermm = 1000.0
    ax.pos_steps = 0.0  # 0 deg

    coupling = {"galil_host": host, "galil_port": port, "axis": axis_letter}
    radius_px = 50.0
    backend = _SimIDSBackend(
        width=200, height=200, rgb=False, rotation_coupling=coupling, feature_radius_px=radius_px
    )
    # same defaults _SimIDSBackend uses internally (speckle_count=60, speckle_seed=42)
    expected_grains = generate_speckle_grains(
        n_speckles=60,
        radius_range=(0.3 * radius_px, radius_px),
        sigma_range=(2.0, 5.0),
        amplitude_range=(120.0, 200.0),
        seed=42,
    )

    frame0 = backend.get_image_data()
    assert np.array_equal(
        frame0,
        make_speckle_test_pattern(200, 200, rgb=False, angle_deg=0.0, grains=expected_grains),
    )

    # move the simulated axis and confirm the frame is regenerated fresh, not cached
    ax.pos_steps = 90.0 * ax.stppermm
    frame90 = backend.get_image_data()
    assert np.array_equal(
        frame90,
        make_speckle_test_pattern(200, 200, rgb=False, angle_deg=90.0, grains=expected_grains),
    )
    assert not np.array_equal(frame0, frame90)


def test_generate_speckle_grains_reproducible():
    """Same seed/params -> identical grains, so the speckle field is fixed
    across frames (only the live angle changes where each grain is drawn)."""
    kwargs = dict(
        n_speckles=20,
        radius_range=(10.0, 50.0),
        sigma_range=(2.0, 4.0),
        amplitude_range=(100.0, 200.0),
        seed=7,
    )
    assert generate_speckle_grains(**kwargs) == generate_speckle_grains(**kwargs)


def test_make_speckle_test_pattern_rotates_rigidly():
    """A single grain's rendered peak moves by exactly the change in
    angle_deg -- the whole speckle field rotates as one rigid body."""
    height, width = 200, 300
    grains = [(60.0, 40.0, 3.0, 200.0)]  # radius_px=60, base_angle_deg=40 (avoids axes)
    cy, cx = height / 2.0, width / 2.0

    def peak_position(angle_deg):
        frame = make_speckle_test_pattern(
            height, width, rgb=False, angle_deg=angle_deg, grains=grains
        )
        masked = frame.astype(np.int32)
        masked[int(cy) - 2 : int(cy) + 3, :] = 0
        masked[:, int(cx) - 2 : int(cx) + 3] = 0
        return np.unravel_index(np.argmax(masked), masked.shape)

    for extra_angle in (0.0, 90.0):
        py, px = peak_position(extra_angle)
        total_angle_rad = np.deg2rad(40.0 + extra_angle)
        expected_y = cy - 60.0 * np.sin(total_angle_rad)
        expected_x = cx + 60.0 * np.cos(total_angle_rad)
        assert abs(py - expected_y) <= 2
        assert abs(px - expected_x) <= 2


def test_sim_ids_camera_threads_rotation_coupling_through():
    """SimIDSCamera's sim_rotation_coupling/sim_feature_radius_px reach _SimIDSBackend."""
    host, port, axis_letter = "test-galil-2.example", 9100, "C"
    galil = SimStateRegistry.get(SimGalilState, host, port)
    ax = galil.axis(ord(axis_letter.lower()) - 97)
    ax.stppermm = 1000.0
    ax.pos_steps = 0.0

    camera = SimIDSCamera(
        name="test_sim_cam",
        camera_id=1,
        prefix="test:",
        scan_info=None,
        m_n_colormode=1,
        bits_per_pixel=24,
        live_mode=False,
        sim_shape=(64, 64),
        sim_rotation_coupling={"galil_host": host, "galil_port": port, "axis": axis_letter},
        sim_feature_radius_px=20.0,
    )

    frame0 = camera.cam.get_image_data()
    ax.pos_steps = 180.0 * ax.stppermm
    frame180 = camera.cam.get_image_data()
    assert not np.array_equal(frame0, frame180)


def test_sim_ids_backend_axis_offset_defaults_to_zero():
    """Without an explicit axis_offset_range_px, the rotation axis stays at
    image center -- unchanged behavior for existing rotation_coupling users."""
    host, port, axis_letter = "test-galil-4.example", 9300, "C"
    galil = SimStateRegistry.get(SimGalilState, host, port)
    ax = galil.axis(ord(axis_letter.lower()) - 97)
    ax.stppermm = 1000.0
    ax.pos_steps = 0.0

    coupling = {"galil_host": host, "galil_port": port, "axis": axis_letter}
    backend = _SimIDSBackend(width=200, height=200, rgb=False, rotation_coupling=coupling)
    assert backend._axis_offset == (0.0, 0.0)


def test_sim_ids_backend_axis_offset_reproducible_with_seed():
    """A given axis_offset_seed always produces the same offset, and the
    rendered frame matches make_speckle_test_pattern called directly with
    that same offset."""
    host, port, axis_letter = "test-galil-3.example", 9200, "C"
    galil = SimStateRegistry.get(SimGalilState, host, port)
    ax = galil.axis(ord(axis_letter.lower()) - 97)
    ax.stppermm = 1000.0
    ax.pos_steps = 0.0

    coupling = {"galil_host": host, "galil_port": port, "axis": axis_letter}
    radius_px = 50.0
    backend = _SimIDSBackend(
        width=200,
        height=200,
        rgb=False,
        rotation_coupling=coupling,
        feature_radius_px=radius_px,
        axis_offset_range_px=30.0,
        axis_offset_seed=99,
    )
    assert backend._axis_offset != (0.0, 0.0)
    assert all(abs(v) <= 30.0 for v in backend._axis_offset)

    # same seed -> same offset, reconstructed independently
    backend2 = _SimIDSBackend(
        width=200,
        height=200,
        rgb=False,
        rotation_coupling=coupling,
        feature_radius_px=radius_px,
        axis_offset_range_px=30.0,
        axis_offset_seed=99,
    )
    assert backend._axis_offset == backend2._axis_offset

    expected_grains = generate_speckle_grains(
        n_speckles=60,
        radius_range=(0.3 * radius_px, radius_px),
        sigma_range=(2.0, 5.0),
        amplitude_range=(120.0, 200.0),
        seed=42,
    )
    frame = backend.get_image_data()
    expected = make_speckle_test_pattern(
        200,
        200,
        rgb=False,
        angle_deg=0.0,
        grains=expected_grains,
        axis_offset=backend._axis_offset,
    )
    assert np.array_equal(frame, expected)


def test_make_speckle_test_pattern_axis_offset_shifts_rotation_center():
    """axis_offset shifts where the grains orbit, away from image center."""
    height, width = 200, 300
    grains = [(0.0, 0.0, 3.0, 200.0)]  # radius=0 -> grain sits exactly on the axis
    axis_offset = (20.0, -10.0)
    frame = make_speckle_test_pattern(
        height, width, rgb=False, angle_deg=0.0, grains=grains, axis_offset=axis_offset
    )
    cy, cx = height / 2.0, width / 2.0
    py, px = np.unravel_index(np.argmax(frame), frame.shape)
    assert abs(py - (cy + axis_offset[1])) <= 2
    assert abs(px - (cx + axis_offset[0])) <= 2