import numpy as np

from orisys import Sensor
from orisys.get_depth import Contact3D
from orisys.hhd import nHHD


def _circle_mask(shape, radius_ratio=0.35):
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w]
    cy = (h - 1) / 2.0
    cx = (w - 1) / 2.0
    radius = min(h, w) * radius_ratio
    return ((yy - cy) ** 2 + (xx - cx) ** 2) <= radius ** 2


def _radial_flow(shape, scale=0.05):
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w]
    cy = (h - 1) / 2.0
    cx = (w - 1) / 2.0
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[:, :, 0] = (xx - cx) * scale
    flow[:, :, 1] = (yy - cy) * scale
    return flow


def test_masked_hhd_zeroes_invalid_region():
    shape = (48, 56)
    mask = _circle_mask(shape)
    flow = _radial_flow(shape, scale=0.03)
    flow[~mask] = 0

    solver = nHHD(cuda=False, grid=shape, spacings=(1.0, 1.0), poisson_backend="cpu")
    _, forces = solver.get_force(flow, valid_mask=mask)

    vnormal = solver.get_normal()
    vshear = solver.get_shear()

    assert np.isfinite(forces).all()
    assert np.allclose(vnormal[~mask], 0.0, atol=1e-5)
    assert np.allclose(vshear[~mask], 0.0, atol=1e-5)


def test_masked_hhd_ignores_invalid_background_spike():
    shape = (48, 56)
    mask = _circle_mask(shape)
    base = _radial_flow(shape, scale=0.02)
    base[~mask] = 0

    contaminated = base.copy()
    contaminated[~mask, 0] = 25.0
    contaminated[~mask, 1] = -15.0

    solver = nHHD(cuda=False, grid=shape, spacings=(1.0, 1.0), poisson_backend="cpu")
    _, forces_base = solver.get_force(base, valid_mask=mask)
    _, forces_contaminated = solver.get_force(contaminated, valid_mask=mask)

    assert np.allclose(forces_base, forces_contaminated, atol=1e-3, rtol=1e-3)


def test_sensor_detect_motion_ignores_invalid_region():
    sensor = Sensor(
        vid_src="external",
        cuda=False,
        isstitch=False,
        verbose=False,
        config_name="dd02-ov",
        mirror=False,
        backend="cpu",
    )
    try:
        # get_img() resizes the frame to the config's camera.resolution_stitch
        # and stores valid_mask at that same resolution, so the flow buffer must
        # be sized from sensor.valid_mask rather than from the raw frame.
        frame = np.zeros((728, 672, 3), dtype=np.uint8)
        mask = np.zeros((728, 672), dtype=bool)
        mask[180:540, 160:512] = True
        sensor.get_img(frame=frame, already_processed=True, valid_mask=mask)

        flow_shape = sensor.valid_mask.shape
        assert sensor.valid_mask.any() and not sensor.valid_mask.all()
        sensor.flow = np.zeros((*flow_shape, 2), dtype=np.float32)
        sensor.flow[~sensor.valid_mask] = 10.0
        assert not sensor.detect_motion(threshold=1.0)

        sensor.flow[:] = 0.0
        sensor.flow[sensor.valid_mask] = 10.0
        assert sensor.detect_motion(threshold=1.0)
    finally:
        sensor.disconnect()


def test_contact_centroid_respects_domain_mask():
    config_contact = {
        "min_contour_area": 25,
        "flow_threshold": 0.01,
        "div_threshold": 0.0,
        "expansion_threshold": 0.0,
        "kernel_size": [3, 3],
        "denoise_sigma": 0.0,
    }
    shape = (96, 96)
    mask = _circle_mask(shape, radius_ratio=0.28)
    flow = _radial_flow(shape, scale=0.05)
    flow[~mask] = 0.0
    flow[:, :8, 0] = 100.0
    flow[:, :8, 1] = 100.0

    contact = Contact3D(config_contact=config_contact, verbose=False)
    result = contact.compute_foe(flow, valid_mask=mask)

    assert result["centroid_found"]
    cx, cy = result["centroid"]
    assert mask[cy, cx]