import json

import numpy as np

from iiwa_genesis.camera import D455, Frame, calibration, optical_pose, optical_rays, process_depth
from iiwa_genesis.config import ROOT, load_config


def camera_config():
    return load_config(ROOT / "configs/scene.yaml")["cameras"]["overhead"]


def test_downward_pose():
    cfg = camera_config()
    pose = optical_pose(cfg["position"], cfg["lookat"], cfg["up"])
    np.testing.assert_allclose(pose[:3, :3].T @ pose[:3, :3], np.eye(3))
    np.testing.assert_allclose(pose[:3, 2], [0, 0, -1])
    assert np.linalg.det(pose[:3, :3]) > 0.999


def test_plane_range_converts_to_constant_z():
    cfg = camera_config()
    distances = 2 * np.linalg.norm(optical_rays(cfg), axis=-1)
    depth, valid = process_depth(distances, cfg, np.random.default_rng(1))
    np.testing.assert_allclose(depth, 2)
    assert valid.all()
    distances[0, :3] = [np.inf, 0, 100]
    depth, valid = process_depth(distances, cfg, np.random.default_rng(1))
    assert not valid[0, :3].any()
    assert not depth[0, :3].any()


def test_seeded_noise_and_serialization(tmp_path):
    cfg = camera_config()
    cfg["noise"]["enabled"] = True
    distances = 2 * np.linalg.norm(optical_rays(cfg), axis=-1)
    depth, valid = process_depth(distances, cfg, np.random.default_rng(7))
    second, _ = process_depth(distances, cfg, np.random.default_rng(7))
    np.testing.assert_array_equal(depth, second)
    assert np.std(depth) > 0
    camera = object.__new__(D455)
    camera.name, camera.cfg = "overhead", cfg
    frame = Frame(np.zeros((480, 640, 3), dtype=np.uint8), depth, valid,
                  {"intrinsics": calibration(cfg).tolist(), "world_from_optical": np.eye(4).tolist()})
    camera.save(frame, tmp_path, 0, pointcloud=True)
    data = np.load(tmp_path / "overhead_000000.npz")
    assert data["depth"].shape == (480, 640)
    assert data["xyz_optical"].shape == (480, 640, 3)
    assert json.loads((tmp_path / "overhead_000000.json").read_text()) == frame.metadata