"""Track a full grid across the synthetic crack with
`locate_warp_propagated`, and map each point's ZNSSD. Two cases: a 4 px
offset with a grid column on the crack, and a 4.3 px offset with the
grid shifted 2 px so no column lines up with it. Prints one table and
saves the off-grid case's ZNSSD and v maps.
"""

import matplotlib.pyplot as plt
import numpy as np

import dictk
from dictk.grid import generate, locate_warp, locate_warp_propagated
from dictk.image import PixelCoordinate, combine, crack_dislocation

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

WIDTH = HEIGHT = 300
KERNEL_MARGIN = 13
SEARCH_MARGIN = 25
GATE = 0.3
CRACK_X = 149.5
CASES = [("on-grid", 4.0, 45), ("off-grid", 4.3, 47)]

reference_image = combine(
    a=dictk.rosta(width=WIDTH, height=HEIGHT, density=0.5),
    b=dictk.astronaut(width=WIDTH, height=HEIGHT),
)
margins = dict(
    kernel_margin_width=KERNEL_MARGIN,
    kernel_margin_height=KERNEL_MARGIN,
    search_margin_width=SEARCH_MARGIN,
    search_margin_height=SEARCH_MARGIN,
)

rows = {}
for name, offset, origin_x in CASES:
    current_image = crack_dislocation(arr=reference_image, offset=offset)
    points = generate(
        origin=PixelCoordinate(x=origin_x, y=45),
        count_x=43,
        count_y=43,
        spacing_x=5,
        spacing_y=5,
    )
    x = np.array([p.x for p in points], dtype=np.float64)
    y = np.array([p.y for p in points], dtype=np.float64)
    true_y = y + np.where(x < CRACK_X, offset, -offset)
    found, znssds = locate_warp_propagated(
        reference_image=reference_image,
        current_image=current_image,
        reference_points=points,
        **margins,
    )
    rigid_seeded = locate_warp(
        reference_image=reference_image,
        current_image=current_image,
        reference_points=points,
        **margins,
    )
    errors = np.hypot([f.x for f in found] - x, [f.y for f in found] - true_y)
    rigid_errors = np.hypot(
        [f.x for f in rigid_seeded] - x, [f.y for f in rigid_seeded] - true_y
    )
    # A kernel spans x - 13 to x + 13, so it straddles the crack when it
    # holds pixels from both columns 149 and 150.
    straddling = np.abs(x - CRACK_X) < KERNEL_MARGIN + 1
    accepted = znssds <= GATE
    columns = np.unique(x)
    medians = [np.median(znssds[x == column]) for column in columns]
    peak_column = columns[np.argmax(medians)]
    rows[name] = [
        f"{offset} px",
        f"{(~straddling).sum()}",
        f"{znssds[~straddling].max():.3f}",
        f"{errors[~straddling].max():.3f} px",
        f"{straddling.sum()}",
        f"{(straddling & ~accepted).sum()}",
        f"{np.isinf(znssds).sum()}",
        f"{(straddling & accepted).sum()}",
        f"{errors[straddling & accepted].max():.3f} px",
        f"{(accepted & ~(errors <= 1.0)).sum()}",
        f"{(rigid_errors > 1.0).sum()}",
        f"{peak_column:.0f} ({abs(peak_column - CRACK_X):.1f} px)",
    ]
    if name == "off-grid":
        map_case = (x, y, znssds, np.array([f.y for f in found]) - y, offset)

labels = [
    "Crack offset",
    "Points away from the crack",
    "... largest ZNSSD",
    "... largest error",
    "Points straddling the crack",
    "... rejected by the gate",
    "... of those, never fitted",
    "... accepted",
    "... largest error when accepted",
    "Accepted points more than 1 px off",
    "`locate_warp` points more than 1 px off",
    "Column with the highest median ZNSSD",
]
print("| | " + " | ".join(name for name, _, _ in CASES) + " |")
print("|---|" + "---|" * len(CASES))
for i, label in enumerate(labels):
    print(f"| {label} | " + " | ".join(rows[name][i] for name, _, _ in CASES) + " |")

x, y, znssds, v, offset = map_case
accepted = znssds <= GATE
fitted = np.isfinite(znssds)
fig, (left, right) = plt.subplots(1, 2, figsize=(10.0, 4.4), constrained_layout=True)
for axis in (left, right):
    axis.imshow(reference_image, cmap="gray", vmin=0, vmax=255)
    axis.axvline(CRACK_X, color="white", linestyle="--", linewidth=0.8)
    axis.set_xlabel("x (pixels)")
    axis.set_ylabel("y (pixels)")
shown = left.scatter(
    x[fitted],
    y[fitted],
    c=np.minimum(znssds[fitted], 1.0),
    s=20,
    marker="s",
    linewidths=0,
    cmap="viridis",
    vmin=0.0,
    vmax=1.0,
)
left.scatter(x[~fitted], y[~fitted], s=20, marker="s", linewidths=0, color="tab:red")
fig.colorbar(shown, ax=left, label="ZNSSD (clipped at 1)")
left.set_title("ZNSSD (red: never fitted)")
shown = right.scatter(
    x[accepted],
    y[accepted],
    c=v[accepted],
    s=20,
    marker="s",
    linewidths=0,
    cmap="coolwarm",
    vmin=-offset,
    vmax=offset,
)
fig.colorbar(shown, ax=right, label="v (pixels)")
right.set_title(f"v, accepted points only (ZNSSD ≤ {GATE})")
fig.savefig("discontinuity_detection_synthetic.png", dpi=300)

"""Track a full grid across the real crack image pair with
`locate_warp_propagated`, at two kernel sizes, and map each point's
ZNSSD. Fits a line to the band of rejected points, checks each accepted
point against its neighbors on the same side, and measures the
displacement across a fainter streak parallel to the crack. Prints two
tables and saves the 27 px kernel's maps.
"""

import matplotlib.pyplot as plt
import numpy as np

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

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

SPACING = 5
GATE = 0.3
NEIGHBOR_REACH = 10
# The largest rigid motion is 37 px. A 50 px search margin covers it
# with room to spare; 65 px matches Experimental Dislocation's.
MARGINS = [(13, 50), (25, 65)]
# A 100 px stretch along the crack, near the middle of the streak.
ALONG_RANGE = (250, 350)
STREAK_BANDS = [(-240, -200), (-200, -185), (-185, -155), (-155, -140), (-140, -100)]

reference_image = read(path="experimental_dislocation_reference.tiff")
current_image = read(path="experimental_dislocation_current.tiff")
size = reference_image.shape[0]

rows = {}
for kernel_margin, search_margin in MARGINS:
    origin = search_margin + SPACING
    count = (size - 2 * origin) // SPACING + 1
    points = generate(
        origin=PixelCoordinate(x=origin, y=origin),
        count_x=count,
        count_y=count,
        spacing_x=SPACING,
        spacing_y=SPACING,
    )
    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=search_margin,
        search_margin_height=search_margin,
    )
    x = np.array([p.x for p in points], dtype=np.float64)
    y = np.array([p.y for p in points], dtype=np.float64)
    u = np.array([f.x for f in found]) - x
    v = np.array([f.y for f in found]) - y
    accepted = znssds <= GATE
    fitted = np.isfinite(znssds)

    # Each row's band center is the mean x of its rejected points, after
    # dropping any more than 40 px from the row's median.
    centers, widths = [], []
    for row in np.unique(y):
        xs = x[(y == row) & ~accepted]
        if len(xs) >= 2:
            xs = xs[np.abs(xs - np.median(xs)) < 40]
            centers.append((row, xs.mean()))
            widths.append(xs.max() - xs.min() + SPACING)
    band_y, band_x = np.array(centers).T
    slope, intercept = np.polyfit(band_y, band_x, 1)
    scatter = np.std(band_x - (slope * band_y + intercept))
    cosine = np.cos(np.arctan(slope))
    across = (x - (slope * y + intercept)) * cosine
    along = (y + slope * (x - intercept)) * cosine
    near = np.abs(across) <= 2 * kernel_margin + 10

    side = np.sign(across)
    deviations = []
    for i in np.flatnonzero(accepted):
        neighbors = (
            accepted
            & (side == side[i])
            & (np.abs(x - x[i]) <= NEIGHBOR_REACH)
            & (np.abs(y - y[i]) <= NEIGHBOR_REACH)
        )
        neighbors[i] = False
        if neighbors.sum() >= 3:
            deviations.append(
                max(
                    abs(u[i] - np.median(u[neighbors])),
                    abs(v[i] - np.median(v[neighbors])),
                )
            )
    deviations = np.array(deviations)
    lower, upper = accepted & (side < 0), accepted & (side > 0)

    rows[kernel_margin] = [
        f"{2 * kernel_margin + 1} px",
        f"{search_margin} px",
        f"{len(points)}",
        f"{accepted.sum()} / {(fitted & ~accepted).sum()} / {(~fitted).sum()}",
        f"{np.median(znssds[fitted]):.3f}",
        f"{(~accepted & near).sum()} / {(~accepted & ~near).sum()}",
        f"x = {slope:.3f} y + {intercept:.1f}",
        f"{np.degrees(np.arctan(slope)):.1f}°",
        f"{scatter:.2f} px",
        f"{np.median(widths):.0f} px",
        f"{(deviations > 1.0).sum()} of {accepted.sum()} ({deviations.max():.2f} px)",
        f"{np.median(u[lower]):.2f} / {np.median(v[lower]):.2f} px",
        f"{np.median(u[upper]):.2f} / {np.median(v[upper]):.2f} px",
    ]
    if kernel_margin == 13:
        map_case = (x, y, znssds, u, v, slope, intercept)
        streak = []
        for low, high in STREAK_BANDS:
            band = (
                accepted
                & (across >= low)
                & (across < high)
                & (along > ALONG_RANGE[0])
                & (along < ALONG_RANGE[1])
            )
            streak.append(
                f"| {low} to {high} px | {band.sum()} | {np.median(znssds[band]):.3f} "
                f"| {np.median(u[band]):.2f} px | {np.median(v[band]):.2f} px |"
            )

labels = [
    "Kernel width",
    "Search margin",
    "Points tracked",
    "Accepted / rejected / never fitted",
    "Median ZNSSD of fitted points",
    "Rejected near the band / farther away",
    "Line fitted to the band",
    "Band angle from vertical",
    "RMS scatter of band centers about the line",
    "Median band width",
    "Accepted points more than 1 px from same-side neighbors (largest)",
    "Lower-left side, median u / v",
    "Upper-right side, median u / v",
]
print("| | " + " | ".join(f"`kernel_margin={k}`" for k, _ in MARGINS) + " |")
print("|---|" + "---|" * len(MARGINS))
for i, label in enumerate(labels):
    print(f"| {label} | " + " | ".join(rows[k][i] for k, _ in MARGINS) + " |")
print()
print("| Distance from the crack line | Points | Median ZNSSD | Median u | Median v |")
print("|---|---|---|---|---|")
print("\n".join(streak))

x, y, znssds, u, v, slope, intercept = map_case
accepted = znssds <= GATE
fitted = np.isfinite(znssds)
line_y = np.array([y.min(), y.max()])
fig, axes = plt.subplots(2, 2, figsize=(10.5, 9.0), constrained_layout=True)
(current_axis, znssd_axis), (u_axis, v_axis) = axes
current_axis.imshow(current_image, cmap="gray", vmin=0, vmax=255)
current_axis.set_title("Current image")
for axis in (znssd_axis, u_axis, v_axis):
    axis.imshow(reference_image, cmap="gray", vmin=0, vmax=255)
shown = znssd_axis.scatter(
    x[fitted],
    y[fitted],
    c=np.minimum(znssds[fitted], 1.0),
    s=8,
    marker="s",
    linewidths=0,
    cmap="viridis",
    vmin=0.0,
    vmax=1.0,
)
znssd_axis.scatter(
    x[~fitted], y[~fitted], s=8, marker="s", linewidths=0, color="tab:red"
)
znssd_axis.plot(
    slope * line_y + intercept, line_y, color="white", linestyle="--", linewidth=0.8
)
fig.colorbar(shown, ax=znssd_axis, label="ZNSSD (clipped at 1)")
znssd_axis.set_title("ZNSSD (red: never fitted)")
for axis, values, name in ((u_axis, u, "u"), (v_axis, v, "v")):
    shown = axis.scatter(
        x[accepted],
        y[accepted],
        c=values[accepted],
        s=8,
        marker="s",
        linewidths=0,
        cmap="coolwarm",
    )
    fig.colorbar(shown, ax=axis, label=f"{name} (pixels)")
    axis.set_title(f"{name}, accepted points only (ZNSSD ≤ {GATE})")
for axis in axes.flat:
    axis.set_xlabel("x (pixels)")
    axis.set_ylabel("y (pixels)")
    axis.set_xlim(0, size)
    axis.set_ylim(size, 0)
fig.savefig("discontinuity_detection_experimental.png", dpi=300)