"""Can one bad region spread? Covers a 40x40 px patch of the current
image with unrelated speckle, under a 15% stretch with a 0.15 shear,
and tracks the 17x17 grid with `locate_warp_propagated` with its ZNSSD
gate off (`znssd_limit=4`, the largest ZNSSD) and on (the default 0.3).
"""

import numpy as np
from scipy.ndimage import affine_transform

from dictk.grid import generate, locate_warp_propagated
from dictk.image import PixelCoordinate, read
from dictk.rosta import rosta

CENTER = np.array([149.5, 149.5])
SHIFT = np.array([2.3, 1.6])
F = np.array([[1.15, 0.15], [0.0, 1.0]])
PATCH = 40
KERNEL_MARGIN = 13
GATE = 0.3

reference_image = read(path="astronaut0.png").astype(np.float64)
points = generate(
    origin=PixelCoordinate(x=110, y=110),
    count_x=17,
    count_y=17,
    spacing_x=5,
    spacing_y=5,
)
reference = np.array([[p.x, p.y] for p in points], dtype=np.float64)
inverse = np.linalg.inv(F)[::-1, ::-1]
offset = CENTER[::-1] - inverse @ (CENTER[::-1] + SHIFT[::-1])
deformed = affine_transform(reference_image, inverse, offset=offset, order=5)
true = CENTER + (reference - CENTER) @ F.T + SHIFT
speckle = rosta(width=PATCH, height=PATCH, density=0.5, random_seed=7)
speckle = speckle.astype(np.float64) * 255.0 / speckle.max()

print(
    "| Patch corner | Gate | Points outside | Bad outside | Bad inside "
    "| Never fitted | Bad but passing the gate | Right but failing the gate |"
)
print("|---|---|---|---|---|---|---|---|")
for corner in [(160, 160), (130, 130)]:
    x0, y0 = corner
    current_image = deformed.copy()
    current_image[y0 : y0 + PATCH, x0 : x0 + PATCH] = speckle
    # A point is inside when its current-image kernel overlaps the patch.
    reach = PATCH / 2 + KERNEL_MARGIN + 0.5
    middle = np.array([x0, y0]) + (PATCH - 1) / 2
    inside = np.all(np.abs(true - middle) < reach, axis=1)
    for gate, limit in [("off", 4.0), ("on", GATE)]:
        found, znssds = locate_warp_propagated(
            reference_image=reference_image,
            current_image=current_image,
            reference_points=points,
            kernel_margin_width=KERNEL_MARGIN,
            kernel_margin_height=KERNEL_MARGIN,
            search_margin_width=25,
            search_margin_height=25,
            znssd_limit=limit,
        )
        errors = np.linalg.norm(np.array([[f.x, f.y] for f in found]) - true, axis=1)
        # A never-fitted point has a nan position, so it counts as bad.
        bad = ~(errors <= 1.0)
        print(
            f"| ({x0}, {y0}) | {gate} | {(~inside).sum()} | {bad[~inside].sum()} "
            f"| {bad[inside].sum()} of {inside.sum()} | {np.isinf(znssds).sum()} "
            f"| {(bad & (znssds <= GATE)).sum()} "
            f"| {(~bad & (znssds > GATE)).sum()} |"
        )

"""`locate_warp` against `locate_warp_propagated` on a 17x17 grid at
5 px spacing, across the sweep's 49 stretch and shear combinations.

Each current image deforms about the image center, then shifts by
(2.3, 1.6) px. `locate_warp` gets each point's true position, rounded,
as its search center, so its search area always contains the true
position and only the rigid seed can fail it. `locate_warp_propagated`
gets no search centers.
"""

import numpy as np
from scipy.ndimage import affine_transform

from dictk.grid import generate, locate_warp, locate_warp_propagated
from dictk.image import PixelCoordinate, read

STRETCHES = np.round(np.arange(0.0, 0.31, 0.05), 2)
SHEARS = np.round(np.arange(0.0, 0.31, 0.05), 2)
MISS_LIMIT = 1.0
CENTER = np.array([149.5, 149.5])
SHIFT = np.array([2.3, 1.6])

reference_image = read(path="astronaut0.png").astype(np.float64)
points = generate(
    origin=PixelCoordinate(x=110, y=110),
    count_x=17,
    count_y=17,
    spacing_x=5,
    spacing_y=5,
)
reference = np.array([[p.x, p.y] for p in points], dtype=np.float64)
margins = dict(
    kernel_margin_width=13,
    kernel_margin_height=13,
    search_margin_width=25,
    search_margin_height=25,
)

names = ["`locate_warp`", "`locate_warp_propagated`"]
miss_counts = np.zeros((2, len(STRETCHES), len(SHEARS)), dtype=int)
largest = np.zeros(2)
for i, stretch in enumerate(STRETCHES):
    for j, shear in enumerate(SHEARS):
        F = np.array([[1.0 + stretch, shear], [0.0, 1.0]])
        # affine_transform maps each output (row, column) back to its
        # input position, so it takes the inverse deformation, in (y, x)
        # order.
        inverse = np.linalg.inv(F)[::-1, ::-1]
        offset = CENTER[::-1] - inverse @ (CENTER[::-1] + SHIFT[::-1])
        current_image = affine_transform(
            reference_image, inverse, offset=offset, order=5
        )
        true = CENTER + (reference - CENTER) @ F.T + SHIFT
        search_centers = [
            PixelCoordinate(x=int(x), y=int(y)) for x, y in np.round(true)
        ]
        for k, (locate, extra) in enumerate(
            [
                (locate_warp, dict(search_centers=search_centers)),
                (locate_warp_propagated, {}),
            ]
        ):
            found = locate(
                reference_image=reference_image,
                current_image=current_image,
                reference_points=points,
                **margins,
                **extra,
            )
            if locate is locate_warp_propagated:
                found, _ = found
            errors = np.linalg.norm(
                np.array([[f.x, f.y] for f in found]) - true, axis=1
            )
            miss_counts[k, i, j] = (errors > MISS_LIMIT).sum()
            largest[k] = max(largest[k], errors.max())

print(f"{names[0]} misses of more than {MISS_LIMIT:g} px, out of {len(points)} points:")
print()
print("| Stretch \\ Shear | " + " | ".join(f"{g:.2f}" for g in SHEARS) + " |")
print("|---|" + "---|" * len(SHEARS))
for stretch, row in zip(STRETCHES, miss_counts[0]):
    print(f"| {stretch:.0%} | " + " | ".join(str(c) for c in row) + " |")
print()

fits = len(points) * len(STRETCHES) * len(SHEARS)
print(f"| Tracker | Misses, of {fits} | Largest error (px) |")
print("|---|---|---|")
for name, counts, error in zip(names, miss_counts, largest):
    print(f"| {name} | {counts.sum()} | {error:.3f} |")

"""How many misses remain when IC-GN starts from the right stretch and
shear? Reruns the sweep's 1470 fits, seeding each from the true
deformation with its translation rounded to whole pixels, as coarse as
the rigid seed's.
"""

import numpy as np
from scipy.ndimage import affine_transform

from dictk import warp
from dictk.image import PixelCoordinate, read

MARGIN = 13
SAMPLE_COUNT = 30
STRETCHES = np.round(np.arange(0.0, 0.31, 0.05), 2)
SHEARS = np.round(np.arange(0.0, 0.31, 0.05), 2)
MISS_LIMIT = 1.0

reference_image = read(path="astronaut0.png")
reference_gradients = warp._gradients(reference_image)
rng = np.random.default_rng(0)
pivots = rng.integers(40, 261, size=(SAMPLE_COUNT, 2)).astype(np.float64)
shifts = np.round(rng.uniform([4, 2], [8, 6], size=(SAMPLE_COUNT, 2)), 2)

errors = []
for stretch in STRETCHES:
    for shear in SHEARS:
        F = np.array([[1.0 + stretch, shear], [0.0, 1.0]])
        inverse = np.linalg.inv(F)[::-1, ::-1]
        for pivot, shift in zip(pivots, shifts):
            point = PixelCoordinate(x=int(pivot[0]), y=int(pivot[1]))
            offset = pivot[::-1] - inverse @ (pivot[::-1] + shift[::-1])
            current_image = affine_transform(
                reference_image.astype(np.float64), inverse, offset=offset, order=5
            )
            initial_warp = np.eye(3)
            initial_warp[:2, :2] = F
            initial_warp[:2, 2] = np.round(shift)
            fitted, _, _ = warp._refine(
                reference_image=reference_image,
                reference_gradients=reference_gradients,
                current_coefficients=warp._spline_coefficients(current_image),
                reference_point=point,
                initial_warp=initial_warp,
                kernel_margin_width=MARGIN,
                kernel_margin_height=MARGIN,
                max_iterations=50,
                tolerance=1e-6,
            )
            errors.append(np.hypot(*(fitted[:2, 2] - shift)))

errors = np.array(errors)
print(
    f"Across all {errors.size} fits, {(errors > MISS_LIMIT).sum()} land more than "
    f"{MISS_LIMIT:g} px from the true position. The largest error is "
    f"{errors.max():.3f} px."
)

"""Can a wrong fit be detected? Reruns the sweep's 1470 fits, the same
49 stretch and shear combinations at the same 30 points, and records
each fit's final ZNSSD and IC-GN iteration count alongside its error.
"""

import matplotlib.pyplot as plt
import numpy as np
from scipy.ndimage import affine_transform

from dictk import translation, warp
from dictk.image import PixelCoordinate, read

plt.rcParams.update({"font.family": "serif", "mathtext.fontset": "cm"})

MARGIN = 13
SEARCH_MARGIN = 25
SAMPLE_COUNT = 30
MAX_ITERATIONS = 50
STRETCHES = np.round(np.arange(0.0, 0.31, 0.05), 2)
SHEARS = np.round(np.arange(0.0, 0.31, 0.05), 2)
MISS_LIMIT = 1.0

reference_image = read(path="astronaut0.png")
reference_gradients = warp._gradients(reference_image)
rng = np.random.default_rng(0)
pivots = rng.integers(40, 261, size=(SAMPLE_COUNT, 2)).astype(np.float64)
shifts = np.round(rng.uniform([4, 2], [8, 6], size=(SAMPLE_COUNT, 2)), 2)

errors, znssds, iterations, cells = [], [], [], []
for stretch in STRETCHES:
    for shear in SHEARS:
        inverse = np.linalg.inv(np.array([[1.0 + stretch, shear], [0.0, 1.0]]))
        inverse = inverse[::-1, ::-1]
        for pivot, shift in zip(pivots, shifts):
            point = PixelCoordinate(x=int(pivot[0]), y=int(pivot[1]))
            offset = pivot[::-1] - inverse @ (pivot[::-1] + shift[::-1])
            current_image = affine_transform(
                reference_image.astype(np.float64), inverse, offset=offset, order=5
            )
            seed = translation.locate(
                reference_image=reference_image,
                current_image=current_image,
                reference_point=point,
                search_center=point,
                kernel_margin_width=MARGIN,
                kernel_margin_height=MARGIN,
                search_margin_width=SEARCH_MARGIN,
                search_margin_height=SEARCH_MARGIN,
            )
            fitted, znssd, count = warp._refine(
                reference_image=reference_image,
                reference_gradients=reference_gradients,
                current_coefficients=warp._spline_coefficients(current_image),
                reference_point=point,
                initial_warp=warp._translation_warp(reference_point=point, start=seed),
                kernel_margin_width=MARGIN,
                kernel_margin_height=MARGIN,
                max_iterations=MAX_ITERATIONS,
                tolerance=1e-6,
            )
            errors.append(np.hypot(*(fitted[:2, 2] - shift)))
            znssds.append(znssd)
            iterations.append(count)
            cells.append((stretch, shear))

errors, znssds, iterations = map(np.array, (errors, znssds, iterations))
missed = errors > MISS_LIMIT
stalled = iterations == MAX_ITERATIONS

print(
    f"| Fit | Count | ZNSSD range | Median ZNSSD | Median iterations "
    rf"| Stopped at $n_{{\text{{MAX\_ITER}}}} = {MAX_ITERATIONS}$ |"
)
print("|---|---|---|---|---|---|")
for name, group in [("Correct", ~missed), ("Miss", missed)]:
    print(
        f"| {name} | {group.sum()} "
        f"| {znssds[group].min():.4f} to {znssds[group].max():.4f} "
        f"| {np.median(znssds[group]):.4f} "
        f"| {np.median(iterations[group]):.0f} "
        f"| {stalled[group].sum()} |"
    )
print()
print("The fits the iteration count misclassifies:")
print()
print("| Fit | Stretch | Shear | Error (px) | ZNSSD | Iterations |")
print("|---|---|---|---|---|---|")
for name, group in [
    ("Correct", ~missed & stalled),
    ("Miss", missed & ~stalled),
]:
    for k in np.flatnonzero(group):
        stretch, shear = cells[k]
        print(
            f"| {name} | {stretch:.0%} | {shear:.2f} | {errors[k]:.3f} "
            f"| {znssds[k]:.4f} | {iterations[k]} |"
        )

fig, (left, right) = plt.subplots(1, 2, figsize=(10, 4.2), constrained_layout=True)
bins = np.logspace(np.log10(znssds.min()) - 0.1, np.log10(znssds.max()) + 0.1, 60)
for group, color, label in [
    (~missed, "tab:blue", "correct fit"),
    (missed, "tab:orange", "miss"),
]:
    left.hist(znssds[group], bins=bins, color=color, alpha=0.8, label=label)
    right.hist(
        iterations[group],
        bins=np.arange(0.5, MAX_ITERATIONS + 1.5),
        color=color,
        alpha=0.8,
        label=label,
    )
left.set_xscale("log")
left.set_yscale("log")
left.set_xlabel("final ZNSSD")
left.set_ylabel("fits")
left.legend()
right.set_yscale("log")
right.set_xlabel("IC-GN iterations")
right.set_ylabel("fits")
right.legend()
fig.savefig("seed_propagation_quality.png", dpi=300)

"""Every whole-pixel seed on Kernel Warping's 2862 points, under its 2%
stretch: how far each seed lands from the true position, and how far
`locate_warp` lands after IC-GN refines it.
"""

import numpy as np

from dictk import translation
from dictk.grid import generate, locate_warp
from dictk.image import PixelCoordinate, read, stretch

FACTOR_X = 1.02
SEED_LIMIT = 1.0

reference_image = read(path="astronaut0.png")
current_image = stretch(arr=reference_image, factor_x=FACTOR_X)
points = generate(
    origin=PixelCoordinate(x=18, y=16),
    count_x=53,
    count_y=54,
    spacing_x=5,
    spacing_y=5,
)
kwargs = dict(
    reference_image=reference_image,
    current_image=current_image,
    kernel_margin_width=13,
    kernel_margin_height=13,
    search_margin_width=25,
    search_margin_height=25,
)
seeds = [
    translation.locate(reference_point=point, search_center=point, **kwargs)
    for point in points
]
found = locate_warp(reference_points=points, **kwargs)

# `image.stretch` pivots on the image origin, so each true x is the
# reference x times the stretch factor, and y is unchanged.
true_positions = np.array([[point.x * FACTOR_X, point.y] for point in points])
seed_errors = np.linalg.norm(
    np.array([[seed.x, seed.y] for seed in seeds]) - true_positions, axis=1
)
warp_errors = np.linalg.norm(
    np.array([[f.x, f.y] for f in found]) - true_positions, axis=1
)

print(
    f"| Points | Seeds more than {SEED_LIMIT:g} px off | Largest seed error (px) "
    "| Largest `locate_warp` error (px) |"
)
print("|---|---|---|---|")
print(
    f"| {len(points)} | {(seed_errors > SEED_LIMIT).sum()} "
    f"| {seed_errors.max():.2f} | {warp_errors.max():.3f} |"
)

"""What a seed is: `translation.locate`'s whole-pixel match, the rigid
starting warp IC-GN refines. Two of the sweep's 30 points under a 15%
stretch and a 0.15 shear, one with a correct seed (top row) and one
with a defeated seed (bottom row).

Each row picks the point whose seed error sits nearest the median seed
error of its group, so neither example is hand-picked.
"""

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Polygon
from scipy.ndimage import affine_transform

from dictk import translation, warp
from dictk.image import PixelCoordinate, read

plt.rcParams.update({"font.family": "serif", "mathtext.fontset": "cm"})

MARGIN = 13
SEARCH_MARGIN = 25
EDGE = 40
SAMPLE_COUNT = 30
SEED_LIMIT = 1.0
F = np.array([[1.15, 0.15], [0.0, 1.0]])
ZOOM = 30

reference_image = read(path="astronaut0.png")
rng = np.random.default_rng(0)
pivots = rng.integers(EDGE, 261, size=(SAMPLE_COUNT, 2)).astype(np.float64)
shifts = np.round(rng.uniform([4, 2], [8, 6], size=(SAMPLE_COUNT, 2)), 2)
inverse = np.linalg.inv(F)[::-1, ::-1]


def deformed(pivot, shift):
    offset = pivot[::-1] - inverse @ (pivot[::-1] + shift[::-1])
    return affine_transform(
        reference_image.astype(np.float64), inverse, offset=offset, order=5
    )


def tracked(pivot, shift):
    point = PixelCoordinate(x=int(pivot[0]), y=int(pivot[1]))
    current_image = deformed(pivot, shift)
    kwargs = dict(
        reference_image=reference_image,
        current_image=current_image,
        reference_point=point,
        search_center=point,
        kernel_margin_width=MARGIN,
        kernel_margin_height=MARGIN,
        search_margin_width=SEARCH_MARGIN,
        search_margin_height=SEARCH_MARGIN,
    )
    seed = translation.locate(**kwargs)
    return current_image, seed, warp.fit(**kwargs)


seed_errors = np.array(
    [
        np.hypot(seed.x - pivot[0] - shift[0], seed.y - pivot[1] - shift[1])
        for pivot, shift in zip(pivots, shifts)
        for _, seed, _ in [tracked(pivot, shift)]
    ]
)
rows = []
for failed in (False, True):
    group = np.flatnonzero((seed_errors > SEED_LIMIT) == failed)
    rows.append(
        group[np.argmin(np.abs(seed_errors[group] - np.median(seed_errors[group])))]
    )

square = np.array(
    [[-MARGIN, -MARGIN], [MARGIN, -MARGIN], [MARGIN, MARGIN], [-MARGIN, MARGIN]],
    dtype=np.float64,
)
fig, axes = plt.subplots(2, 3, figsize=(13, 8.6), constrained_layout=True)
for (kernel_ax, surface_ax, current_ax), index, label in zip(
    axes, rows, ["correct seed", "defeated seed"]
):
    pivot, shift = pivots[index], shifts[index]
    x, y = int(pivot[0]), int(pivot[1])
    current_image, seed, fitted = tracked(pivot, shift)
    true_position = pivot + shift
    seed_offset = np.array([seed.x, seed.y]) - pivot

    kernel_ax.imshow(reference_image, cmap="gray", vmin=0, vmax=255)
    kernel_ax.add_patch(
        Polygon(
            pivot + square,
            fill=False,
            edgecolor="white",
            linestyle="--",
            linewidth=1.5,
        )
    )
    kernel_ax.plot(x, y, marker="o", color="red", markersize=7, mec="black")
    kernel_ax.set_xlim(x - ZOOM, x + ZOOM)
    kernel_ax.set_ylim(y + ZOOM, y - ZOOM)
    kernel_ax.set_title(f"{label}: reference kernel at ({x}, {y})")

    # locate's own phase correlation, step for step as its
    # skimage.registration.phase_cross_correlation call computes it: the
    # kernel centered in a zero-padded copy the search area's size, the
    # cross-power spectrum normalized to unit magnitude, and the peak of
    # the inverse transform's magnitude. After fftshift, the surface's
    # index minus SEARCH_MARGIN is the displacement from the point.
    kernel = reference_image[y - MARGIN : y + MARGIN, x - MARGIN : x + MARGIN]
    search = current_image[
        y - SEARCH_MARGIN : y + SEARCH_MARGIN, x - SEARCH_MARGIN : x + SEARCH_MARGIN
    ]
    pad = SEARCH_MARGIN - MARGIN
    kernel_padded = np.pad(kernel.astype(np.float64), pad)
    product = np.fft.fft2(search) * np.fft.fft2(kernel_padded).conj()
    product /= np.maximum(np.abs(product), 100 * np.finfo(np.float64).eps)
    surface = np.fft.fftshift(np.abs(np.fft.ifft2(product)))
    peak_row, peak_column = np.unravel_index(np.argmax(surface), surface.shape)
    assert (peak_column - SEARCH_MARGIN, peak_row - SEARCH_MARGIN) == tuple(seed_offset)
    extent = [-SEARCH_MARGIN - 0.5, SEARCH_MARGIN - 0.5] * 2
    surface_ax.imshow(surface, cmap="viridis", extent=extent[:2] + extent[2:][::-1])
    surface_ax.plot(
        *seed_offset, marker="s", markersize=9, mfc="none", mec="orange", mew=2
    )
    surface_ax.plot(*shift, marker="+", color="red", markersize=14, mew=2)
    surface_ax.set_xlabel("$u$ (pixels)")
    surface_ax.set_ylabel("$v$ (pixels)")
    surface_ax.set_title(
        f"phase correlation: seed $(u_0, v_0) = ({seed_offset[0]:.0f}, "
        f"{seed_offset[1]:.0f})$"
    )

    current_ax.imshow(current_image, cmap="gray", vmin=0, vmax=255)
    current_ax.add_patch(
        Polygon(
            pivot + seed_offset + square,
            fill=False,
            edgecolor="orange",
            linestyle="--",
            linewidth=2,
            label="seed: rigid square",
        )
    )
    current_ax.add_patch(
        Polygon(
            pivot + fitted[:2, 2] + square @ fitted[:2, :2].T,
            fill=False,
            edgecolor="orange",
            linewidth=2,
            label="IC-GN's final warp",
        )
    )
    current_ax.plot(*true_position, marker="+", color="red", markersize=14, mew=2)
    error = np.hypot(*(fitted[:2, 2] - shift))
    current_ax.set_xlim(x - ZOOM, x + ZOOM)
    current_ax.set_ylim(y + ZOOM, y - ZOOM)
    current_ax.set_title(f"current image: IC-GN misses by {error:.3f} px")
    current_ax.legend(loc="lower left", fontsize=9, framealpha=0.85)

for ax in axes[:, [0, 2]].ravel():
    ax.set_xlabel("x (pixels)")
    ax.set_ylabel("y (pixels)")

fig.savefig("seed_propagation_seed.png", dpi=300)

"""The sweep's 30 points on `astronaut0` (left), and one point's own
deformation, pivoted on that point, then shifted (right).

Draws the same points and shifts as seed_propagation_baseline.py, from
the same seed. The close-up uses the point whose shift sits nearest the
median shift, so the example is typical, not hand-picked.
"""

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Polygon, Rectangle
from scipy.ndimage import affine_transform

from dictk.image import read

plt.rcParams.update({"font.family": "serif", "mathtext.fontset": "cm"})

MARGIN = 13
EDGE = 40
SAMPLE_COUNT = 30
# 15% stretch with 0.15 shear: inside the sweep, and large enough to see.
F = np.array([[1.15, 0.15], [0.0, 1.0]])
ZOOM = 30

reference_image = read(path="astronaut0.png")
height, width = reference_image.shape
rng = np.random.default_rng(0)
pivots = rng.integers(EDGE, 261, size=(SAMPLE_COUNT, 2)).astype(np.float64)
shifts = np.round(rng.uniform([4, 2], [8, 6], size=(SAMPLE_COUNT, 2)), 2)

typical = int(np.argmin(np.linalg.norm(shifts - np.median(shifts, axis=0), axis=1)))
pivot, shift = pivots[typical], shifts[typical]
inverse = np.linalg.inv(F)[::-1, ::-1]
offset = pivot[::-1] - inverse @ (pivot[::-1] + shift[::-1])
current_image = affine_transform(
    reference_image.astype(np.float64), inverse, offset=offset, order=5
)

square = np.array(
    [[-MARGIN, -MARGIN], [MARGIN, -MARGIN], [MARGIN, MARGIN], [-MARGIN, MARGIN]],
    dtype=np.float64,
)
deformed = square @ F.T

fig, (left, right) = plt.subplots(1, 2, figsize=(10, 5.2), constrained_layout=True)

left.imshow(reference_image, cmap="gray", vmin=0, vmax=255)
left.add_patch(
    Rectangle(
        (EDGE, EDGE),
        width - 1 - 2 * EDGE,
        height - 1 - 2 * EDGE,
        fill=False,
        edgecolor="white",
        linestyle="--",
        linewidth=1.5,
    )
)
others = np.delete(pivots, typical, axis=0)
left.scatter(*others.T, s=30, color="orange", edgecolor="black", linewidth=0.5)
left.scatter(*pivot, s=60, color="red", edgecolor="black", linewidth=0.5)
left.set_title(f"{SAMPLE_COUNT} reference points, at least {EDGE} px from each edge")
left.set_xlabel("x (pixels)")
left.set_ylabel("y (pixels)")

right.imshow(current_image, cmap="gray", vmin=0, vmax=255)
right.add_patch(
    Polygon(
        pivot + square,
        fill=False,
        edgecolor="white",
        linestyle="--",
        linewidth=1.5,
        label="reference kernel",
    )
)
right.add_patch(
    Polygon(
        pivot + deformed,
        fill=False,
        edgecolor="orange",
        linestyle=":",
        linewidth=2,
        label="deformed about the point",
    )
)
right.add_patch(
    Polygon(
        pivot + shift + deformed,
        fill=False,
        edgecolor="orange",
        linewidth=2,
        label="then shifted",
    )
)
right.annotate(
    "",
    xy=pivot + shift,
    xytext=pivot,
    arrowprops=dict(arrowstyle="-|>", color="orange", linewidth=2),
)
right.plot(*pivot, marker="o", color="red", markersize=7, mec="black")
right.plot(*(pivot + shift), marker="+", color="red", markersize=14, mew=2)
right.set_xlim(pivot[0] - ZOOM, pivot[0] + ZOOM + shift[0])
right.set_ylim(pivot[1] + ZOOM + shift[1], pivot[1] - ZOOM)
right.set_title(
    f"point ({pivot[0]:.0f}, {pivot[1]:.0f}), shift "
    f"({shift[0]:.2f}, {shift[1]:.2f}) px\n15% stretch, 0.15 shear"
)
right.set_xlabel("x (pixels)")
right.set_ylabel("y (pixels)")
right.legend(loc="lower left", fontsize=9, framealpha=0.85)

fig.savefig("seed_propagation_points.png", dpi=300)

"""Where `locate_warp`'s rigid whole-pixel seed breaks: a sweep of
stretch and shear, each cell tracked at the same 30 random points, each
under its own random fractional shift.

Random points, not one, because a single point's image content can make
some deformations look easier than smaller ones. Each deformation
pivots on its tracked point, so that point's own displacement is just
the shift. The shift never exceeds 10 px, inside
the 12 px a 25 px search margin leaves around a 13 px kernel margin.
The search area therefore always contains the true position, and any
failure comes from the rigid kernel's mismatch alone.
"""

import numpy as np
from scipy.ndimage import affine_transform

from dictk import translation, warp
from dictk.image import PixelCoordinate, read

MARGIN = 13
SEARCH_MARGIN = 25
SAMPLE_COUNT = 30
STRETCHES = np.round(np.arange(0.0, 0.31, 0.05), 2)
SHEARS = np.round(np.arange(0.0, 0.31, 0.05), 2)
# A correct whole-pixel match lies within sqrt(0.5^2 + 0.5^2) = 0.71 px
# of the truth. More than 1 px means at least one whole pixel off.
SEED_LIMIT = 1.0
MISS_LIMIT = 1.0

reference_image = read(path="astronaut0.png")
rng = np.random.default_rng(0)
# Points at least 40 px from every edge of the 300x300 image, so the
# 25 px search area plus an 8 px shift stays inside it.
pivots = rng.integers(40, 261, size=(SAMPLE_COUNT, 2)).astype(np.float64)
# The same shift range as Kernel Warping, rounded to 0.01 px so each
# shift prints exactly as used.
shifts = np.round(rng.uniform([4, 2], [8, 6], size=(SAMPLE_COUNT, 2)), 2)


def deformed(F, pivot, shift):
    # affine_transform maps each output (row, column) back to its input
    # position, so it takes the inverse deformation, in (y, x) order.
    inverse = np.linalg.inv(F)[::-1, ::-1]
    offset = pivot[::-1] - inverse @ (pivot[::-1] + shift[::-1])
    return affine_transform(
        reference_image.astype(np.float64), inverse, offset=offset, order=5
    )


seed_errors = np.zeros((len(STRETCHES), len(SHEARS), SAMPLE_COUNT))
warp_errors = np.zeros_like(seed_errors)
for i, stretch in enumerate(STRETCHES):
    for j, shear in enumerate(SHEARS):
        F = np.array([[1.0 + stretch, shear], [0.0, 1.0]])
        for k, (pivot, shift) in enumerate(zip(pivots, shifts)):
            point = PixelCoordinate(x=int(pivot[0]), y=int(pivot[1]))
            current_image = deformed(F, pivot, shift)
            true_position = pivot + shift
            for errors, locate in [
                (seed_errors, translation.locate),
                (warp_errors, warp.locate),
            ]:
                found = locate(
                    reference_image=reference_image,
                    current_image=current_image,
                    reference_point=point,
                    search_center=point,
                    kernel_margin_width=MARGIN,
                    kernel_margin_height=MARGIN,
                    search_margin_width=SEARCH_MARGIN,
                    search_margin_height=SEARCH_MARGIN,
                )
                errors[i, j, k] = np.hypot(
                    found.x - true_position[0], found.y - true_position[1]
                )

seed_counts = (seed_errors > SEED_LIMIT).sum(axis=2)
miss_counts = (warp_errors > MISS_LIMIT).sum(axis=2)


def table(counts, title):
    print(f"{title}, out of {SAMPLE_COUNT} points:")
    print()
    print("| Stretch \\ Shear | " + " | ".join(f"{g:.2f}" for g in SHEARS) + " |")
    print("|---|" + "---|" * len(SHEARS))
    for stretch, row in zip(STRETCHES, counts):
        print(f"| {stretch:.0%} | " + " | ".join(str(c) for c in row) + " |")
    print()


table(seed_counts, f"Seeds more than {SEED_LIMIT:g} px off")
table(miss_counts, f"`locate_warp` misses of more than {MISS_LIMIT:g} px")
converged = warp_errors[warp_errors <= MISS_LIMIT]
missed = warp_errors[warp_errors > MISS_LIMIT]
print(
    f"Across all {warp_errors.size} fits, the largest error at or under "
    f"{MISS_LIMIT:g} px is {converged.max():.3f} px, and the smallest error "
    f"over it is {missed.min():.2f} px."
)
seed_failed = seed_errors > SEED_LIMIT
warp_missed = warp_errors > MISS_LIMIT
print()
print(
    f"Across the sweep, the seed fails {seed_failed.sum()} times and "
    f"`locate_warp` misses {warp_missed.sum()} times. IC-GN rescues "
    f"{(seed_failed & ~warp_missed).sum()} failed seeds. "
    f"{(warp_missed & ~seed_failed).sum()} misses start from a correct seed."
)