#!/usr/bin/env python3

import sys
import numpy as np
import numpysane as nps
import os

testdir = os.path.dirname(os.path.realpath(__file__))

# I import the LOCAL mrcal since that's what I'm testing
sys.path[:0] = (f"{testdir}/..",)
import mrcal
import testutils

import cv2

image = mrcal.load_image(
    f"{testdir}/data/figueroa-overpass-looking-S.0.downsampled.jpg",
    bits_per_pixel=8,
    channels=1,
)

# some made-up homography with scaling, rotating, skewing and translating
H01 = np.array(
    (
        (0.7, 0.3, 1.0),
        (-0.4, 0.95, 2.0),
        (3e-5, 4e-6, 1.0),
    ),
    dtype=np.float32,
)

H10 = np.linalg.inv(H01)

# The feature I'm going to test with. This is the corner of one of the towers
q0 = np.array((294, 159), dtype=np.float32)

# The transformed image. The matcher should end-up reversing this
# transformation, since it will be given the homography.
#
# shape (H,W,2)
image1 = mrcal.transform_image(
    image,
    mrcal.apply_homography(
        H01,
        nps.glue(
            *[
                nps.dummy(arr, -1)
                for arr in np.meshgrid(np.arange(500), np.arange(600))
            ],
            axis=-1,
        ).astype(np.float32),
    ),
)

# I have the source images and the "right" homography and the "right" matching
# pixels coords. Run the matcher, and compare

templatesize = (30, 20)
search_radius = 50

H10_shifted = H10.copy()
H10_shifted[0, 2] += 10.2
H10_shifted[1, 2] -= 20.4

q1_matched, diagnostics = mrcal.match_feature(
    image,
    image1,
    q0,
    H10=H10_shifted,
    search_radius1=50,
    template_size1=templatesize,
    method=cv2.TM_CCOEFF_NORMED,
)
testutils.confirm_equal(
    q1_matched,
    mrcal.apply_homography(H10, q0),
    worstcase=True,
    eps=0.1,
    msg="match_feature(method=TM_CCOEFF_NORMED) reports the correct pixel coordinate",
)

q1_matched, diagnostics = mrcal.match_feature(
    image,
    image1,
    q0,
    H10=H10_shifted,
    search_radius1=50,
    template_size1=templatesize,
    method=cv2.TM_SQDIFF_NORMED,
)
testutils.confirm_equal(
    q1_matched,
    mrcal.apply_homography(H10, q0),
    worstcase=True,
    eps=0.1,
    msg="match_feature(method=TM_SQDIFF_NORMED) reports the correct pixel coordinate",
)

q1_matched, diagnostics = mrcal.match_feature(
    image,
    image1,
    q0,
    H10=H10_shifted,
    search_radius1=1000,
    template_size1=templatesize,
    method=cv2.TM_CCOEFF_NORMED,
)
testutils.confirm_equal(
    q1_matched,
    mrcal.apply_homography(H10, q0),
    worstcase=True,
    eps=0.1,
    msg="out-of-bounds search_radius works ok",
)
templatesize_hw = np.array((templatesize[-1], templatesize[-2]))
testutils.confirm_equal(
    diagnostics["matchoutput_image"].shape,
    image1.shape - templatesize_hw + 1,
    msg="out-of-bounds search radius looks at the whole image",
)

q1_matched, diagnostics = mrcal.match_feature(
    image * 0,
    image1,
    q0,
    H10=H10_shifted,
    search_radius1=50,
    template_size1=templatesize,
    method=cv2.TM_CCOEFF_NORMED,
)
testutils.confirm_equal(q1_matched, None, msg="failing correlation returns None")

try:
    mrcal.match_feature(
        image * 0,
        image1,
        q0,
        H10=H10_shifted,
        search_radius1=50,
        template_size1=(5000, 5000),
        method=cv2.TM_CCOEFF_NORMED,
    )
except:
    testutils.confirm(True, msg="Too-big template size throws an exception")
else:
    testutils.confirm(False, msg="Too-big template size throws an exception")

testutils.finish()
