import cv2
import numpy as np
import yaml

from hermes_vision.calibrate import (board_points, camera_info_yaml, coverage, novel_enough, solve, solve_robust,
                                     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


def test_k3_is_held_at_zero_by_default_but_can_be_fitted():
    obj, views = synthetic_views()
    _, _, d = solve([obj] * len(views), views, SIZE)
    assert d[4] == 0.0
    _, _, d_full = solve([obj] * len(views), views, SIZE, fix_k3=False)
    assert d_full[4] != 0.0


def test_coverage_names_the_regions_the_board_never_reached():
    w, h = SIZE
    centre_only = [np.random.default_rng(1).uniform([w * 0.4, h * 0.4], [w * 0.6, h * 0.6], (54, 2)).astype(np.float32)]
    frac, empty = coverage(centre_only, SIZE)
    assert frac < 0.2 and 'top-left' in empty and 'bottom-right' in empty and 'centre' not in empty
    rng = np.random.default_rng(2)
    everywhere = [rng.uniform([0, 0], [w, h], (200, 2)).astype(np.float32)]
    assert coverage(everywhere, SIZE) == (1.0, [])


def test_blurred_outlier_views_are_dropped_and_the_result_stays_accurate():
    obj, views = synthetic_views(n=40, seed=9)
    rng = np.random.default_rng(5)
    noisy = [v + rng.normal(0, 0.15, v.shape).astype(np.float32) for v in views]       # ordinary jitter
    bad = [v + rng.normal(0, 2.0, v.shape).astype(np.float32) for v in views[:4]]       # four blurred views
    allv = noisy + bad
    rms_all, k_all, _ = solve([obj] * len(allv), allv, SIZE)
    rms, k, d, keep = solve_robust([obj] * len(allv), allv, SIZE)
    assert set(range(len(noisy), len(allv))) <= set(range(len(allv))) - set(keep)       # the four bad ones are gone
    assert len(keep) >= len(noisy) - 2                                                   # the ordinary ones are kept
    assert rms < rms_all and rms < 0.3
    assert abs(k[0, 0] - 1100) / 1100 < 0.02


def test_it_never_drops_below_the_minimum_number_of_views():
    obj, views = synthetic_views(n=14, seed=4)
    rng = np.random.default_rng(1)
    allv = [v + rng.normal(0, 1.0, v.shape).astype(np.float32) * (i % 2) for i, v in enumerate(views)]
    _, _, _, keep = solve_robust([obj] * len(allv), allv, SIZE, min_views=len(allv) - 1)
    assert len(keep) >= len(allv) - 1