"""Behavior checks with independent spherical and room geometry references."""

import numpy as np

from geometry import (camera_bases, merge_distances, panorama_directions,
                      room_reference, sample_panorama, view_rays)
from prototype import fixed_fov_distance


def test_fixed_fov_recovers_shift_and_metric_distance():
    rays = view_rays(64)
    depth = 2 + 0.3 * rays[..., 0]
    points = rays * depth[..., None]
    points[..., 2] -= 0.7
    distance, mask = fixed_fov_distance(points, np.ones((64, 64)), 1.5)
    np.testing.assert_allclose(distance, depth * np.linalg.norm(rays, axis=-1) * 1.5, rtol=1e-5)
    assert mask.all()


def test_longitude_sampling_wraps_without_a_seam():
    directions = panorama_directions(256, 128)
    image = 2 + directions[..., 0]
    rays = np.array(((1, 1e-8, 0), (1, -1e-8, 0)))
    sampled = sample_panorama(image, rays)
    np.testing.assert_allclose(sampled[0], sampled[1], atol=1e-6)
    np.testing.assert_allclose(sampled, 3, atol=0.001)


def test_fusion_removes_per_view_scale_offsets():
    bases = camera_bases()
    scales = np.exp(np.linspace(-0.4, 0.4, 12))
    distances = []
    for scale, basis in zip(scales, bases):
        rays = view_rays(128) @ basis
        rays /= np.linalg.norm(rays, axis=-1, keepdims=True)
        distances.append((3 + 0.4 * rays[..., 0] + 0.2 * rays[..., 2]) * scale)
    merged, mask = merge_distances(distances, np.ones((12, 128, 128), bool), bases, 64, 32)
    rays = panorama_directions(64, 32)
    expected = 3 + 0.4 * rays[..., 0] + 0.2 * rays[..., 2]
    merged *= np.median(expected / merged)
    assert mask.all()
    assert np.mean(np.abs(merged / expected - 1)) < 0.002
    assert np.max(np.abs(merged[:, [0, -1]] / expected[:, [0, -1]] - 1)) < 0.005


def test_room_distance_lands_on_floor():
    rays = np.array(((0.0, 0.0, -1.0), (0.2, -0.3, -1.0)))
    rays /= np.linalg.norm(rays, axis=-1, keepdims=True)
    distance, _ = room_reference(rays)
    np.testing.assert_allclose((rays * distance[:, None])[:, 2], -1.6)


def test_sky_view_returns_invalid_samples():
    distance, mask = fixed_fov_distance(np.ones((32, 32, 3)), np.zeros((32, 32)), 1)
    assert not mask.any()
    assert np.isfinite(distance).all()


def test_room_projection_fusion_preserves_geometry():
    bases = camera_bases()
    distances = [room_reference(view_rays(256) @ basis)[0] for basis in bases]
    merged, mask = merge_distances(distances, np.ones((12, 256, 256), bool), bases, 128, 64)
    expected, _ = room_reference(panorama_directions(128, 64))
    assert mask.all()
    assert np.mean(np.abs(merged / expected - 1)) < 0.005


def test_input_alpha_survives_panorama_export(tmp_path, monkeypatch):
    from PIL import Image
    from prototype import main

    rgba = np.full((16, 32, 4), 255, dtype=np.uint8)
    rgba[:, :16, 3] = 0
    source = tmp_path / "source.png"
    destination = tmp_path / "result"
    Image.fromarray(rgba).save(source)
    monkeypatch.setattr("sys.argv", ["prototype", "--image", str(source), "--output", str(destination),
                                    "--reference-only", "--size", "32", "--width", "32"])
    main()
    with Image.open(destination / "panorama.png") as exported:
        np.testing.assert_array_equal(np.asarray(exported)[..., 3], rgba[..., 3])
    with np.load(destination / "prediction.npz") as data:
        assert np.all(data["alpha"][:, :16] == 0)
        assert np.any(data["alpha"][:, 16:] > 0)


def test_depth_alpha_combines_mask_and_source_alpha_once():
    from prototype import depth_alpha

    image = np.ones((2, 4, 4), dtype=np.float32)
    image[..., 3] = (0, 0.5, 0.75, 1)
    mask = np.full((2, 4), 0.5, dtype=np.float32)
    np.testing.assert_allclose(depth_alpha(image, mask), np.tile((0, 0.25, 0.375, 0.5), (2, 1)))