import cv2
import numpy as np
import yaml

from hermes_vision.calibrate import board_points, camera_info_yaml, novel_enough, solve, view_signature

SIZE = (1296, 972)
K_TRUE = np.array([[1100.0, 0, 650.0], [0, 1090.0, 480.0], [0, 0, 1]])
D_TRUE = np.array([-0.18, 0.05, 0.0005, -0.0003, 0.0])


def synthetic_views(n=14, seed=3):
    rng = np.random.default_rng(seed)
    obj = board_points(9, 6, 0.025)
    views = []
    for _ in range(n):
        rvec = rng.uniform(-0.45, 0.45, 3)
        tvec = np.array([rng.uniform(-0.06, 0.02), rng.uniform(-0.05, 0.01), rng.uniform(0.28, 0.5)])
        pts, _ = cv2.projectPoints(obj, rvec, tvec, K_TRUE, D_TRUE)
        pts = pts.reshape(-1, 2)
        if (pts[:, 0].min() > 20 and pts[:, 0].max() < SIZE[0] - 20 and pts[:, 1].min() > 20
                and pts[:, 1].max() < SIZE[1] - 20):
            views.append(pts.reshape(-1, 1, 2).astype(np.float32))
    return obj, views


def test_solver_recovers_the_true_camera():
    obj, views = synthetic_views()
    assert len(views) >= 8
    rms, k, d = solve([obj] * len(views), views, SIZE)
    assert rms < 0.05                                          # synthetic points are exact
    assert abs(k[0, 0] - 1100) / 1100 < 0.01 and abs(k[1, 1] - 1090) / 1090 < 0.01
    assert abs(k[0, 2] - 650) < 5 and abs(k[1, 2] - 480) < 5
    assert abs(d[0] - D_TRUE[0]) < 0.03


def test_the_yaml_is_the_camera_info_format_and_round_trips():
    obj, views = synthetic_views()
    _, k, d = solve([obj] * len(views), views, SIZE)
    text = camera_info_yaml('cam_1296x972', SIZE, k, d)
    y = yaml.safe_load(text)
    assert (y['image_width'], y['image_height'], y['camera_name']) == (1296, 972, 'cam_1296x972')
    assert y['distortion_model'] == 'plumb_bob' and len(y['distortion_coefficients']['data']) == 5
    assert np.allclose(np.array(y['camera_matrix']['data']).reshape(3, 3), k, rtol=1e-6)
    assert y['projection_matrix']['rows'] == 3 and y['projection_matrix']['cols'] == 4
    assert y['rectification_matrix']['data'] == [1, 0, 0, 0, 1, 0, 0, 0, 1]


def test_board_points_layout():
    p = board_points(9, 6, 0.025)
    assert p.shape == (54, 3) and np.allclose(p[1] - p[0], [0.025, 0, 0]) and np.allclose(p[9] - p[0], [0, 0.025, 0])


def test_novelty_rejects_repeats_and_accepts_a_real_change():
    _, views = synthetic_views()
    sigs = [view_signature(v, SIZE) for v in views]
    assert not novel_enough(sigs[0], [sigs[0]])                # the same pose again
    far = [s for s in sigs[1:] if np.linalg.norm(s - sigs[0]) > 0.06]
    assert far and novel_enough(far[0], [sigs[0]])
    assert novel_enough(sigs[0], [])                            # the first view is always accepted


def test_find_board_detects_a_real_rendered_checkerboard_and_not_a_blank_frame():
    from hermes_vision.calibrate import find_board
    sq, margin = 80, 80
    img = np.full((7 * sq + 2 * margin, 10 * sq + 2 * margin), 255, np.uint8)
    for r in range(7):
        for c in range(10):
            if (r + c) % 2 == 0:
                img[margin + r * sq: margin + (r + 1) * sq, margin + c * sq: margin + (c + 1) * sq] = 0
    bgr = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    corners = find_board(bgr, 9, 6)
    assert corners is not None and corners.reshape(-1, 2).shape == (54, 2)
    assert find_board(np.full_like(bgr, 128), 9, 6) is None