# examples/surf.py
"""Offline, integer-grid U-SURF mechanics on deterministic synthetic images.

Authored from Bay, Tuytelaars, and Van Gool, ECCV 2006, Sections 3--4.2:
https://people.ee.ethz.ch/~surf/eccv06.pdf . No upstream code is copied.
This is a complete upright 64-D descriptor with a separate orientation helper,
not full SURF: subpixel refinement, scale interpolation, rotated descriptor
sampling, octave resampling, and benchmark reproduction are omitted. Coordinates
are (row, column); rectangles are half-open; out-of-image filters are rejected.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path
import time

import numpy as np

FILTER_SIZES = (9, 15, 21, 27)


def _integer(value):
    """Return whether a coordinate or size is an integer (excluding bool)."""
    return isinstance(value, (int, np.integer)) and not isinstance(value, bool)


def _real_array(value):
    """Convert real numeric inputs without silently discarding complex parts."""
    array = np.asarray(value)
    if array.dtype.kind not in "iuf":
        raise ValueError("Expected a real numeric array.")
    return array.astype(np.float64)


def _validate_sat(sat):
    """Validate a finite zero-padded summed-area table."""
    sat = _real_array(sat)
    if sat.ndim != 2 or min(sat.shape) < 2 or not np.isfinite(sat).all():
        raise ValueError(
            "Expected a finite, nonempty, zero-padded integral image."
        )
    if np.any(sat[0]) or np.any(sat[:, 0]):
        raise ValueError("The first row and column must be zero.")
    return sat


# region: integral
def integral_image(image):
    """Return a float64 SAT with a zero row and column, shape [H+1, W+1]."""
    image = _real_array(image)
    if image.ndim != 2 or not image.size or not np.isfinite(image).all():
        raise ValueError("Expected a finite, nonempty 2-D image.")
    sat = np.zeros((image.shape[0] + 1, image.shape[1] + 1))
    sat[1:, 1:] = image.cumsum(axis=0).cumsum(axis=1)
    if not np.isfinite(sat).all():
        raise ValueError("Integral-image accumulation overflowed.")
    return sat


# endregion: integral


# region: rectangles
def _rectangle(sat, y0, x0, y1, x1):
    """Read four corners; callers establish valid half-open bounds."""
    return float(sat[y1, x1] - sat[y0, x1] - sat[y1, x0] + sat[y0, x0])


def rectangle_sum(sat, y0, x0, y1, x1):
    """Sum [y0:y1, x0:x1], accepting borders and empty rectangles."""
    sat = _validate_sat(sat)
    if not all(_integer(v) for v in (y0, x0, y1, x1)):
        raise ValueError("Rectangle coordinates must be integers.")
    if not (0 <= y0 <= y1 < sat.shape[0] and 0 <= x0 <= x1 < sat.shape[1]):
        raise ValueError(
            "Rectangle lies outside the image or has reversed bounds."
        )
    return _rectangle(sat, y0, x0, y1, x1)


# endregion: rectangles


# region: box_hessian
def box_hessian(sat, y, x, size):
    """Return normalized Dxx, Dyy, Dxy, determinant, and trace sign.

    Odd lobes reproduce Figure 1: Dxx has three vertical lobes with weights
    +1,-2,+1 and height 2*lobe-1; Dyy is its transpose. Dxy uses four lobe-sized
    quadrants separated by a zero center row/column. Divide responses by size².
    """
    sat = _validate_sat(sat)
    if not _integer(size) or size < 9 or size % 6 != 3:
        raise ValueError("Filter size must be 9, 15, 21, 27, ... .")
    if not _integer(y) or not _integer(x):
        raise ValueError("Filter centers must be integer coordinates.")
    lobe, border = size // 3, size // 2
    if not (
        border <= y < sat.shape[0] - 1 - border
        and border <= x < sat.shape[1] - 1 - border
    ):
        raise ValueError("The full box filter must fit inside the image.")

    def r(y0, x0, y1, x1):
        return _rectangle(sat, y0, x0, y1, x1)

    dxx = r(y - lobe + 1, x - border, y + lobe, x + border + 1)
    dxx -= 3 * r(y - lobe + 1, x - lobe // 2, y + lobe, x + lobe // 2 + 1)
    dyy = r(y - border, x - lobe + 1, y + border + 1, x + lobe)
    dyy -= 3 * r(y - lobe // 2, x - lobe + 1, y + lobe // 2 + 1, x + lobe)
    dxy = r(y - lobe, x - lobe, y, x) + r(
        y + 1, x + 1, y + lobe + 1, x + lobe + 1
    )
    dxy -= r(y - lobe, x + 1, y, x + lobe + 1) + r(
        y + 1, x - lobe, y + lobe + 1, x
    )
    dxx, dyy, dxy = np.array([dxx, dyy, dxy]) / size**2
    determinant = dxx * dyy - (0.9 * dxy) ** 2
    return (
        float(dxx),
        float(dyy),
        float(dxy),
        float(determinant),
        int(np.sign(dxx + dyy)),
    )


# endregion: box_hessian


# region: scale_maxima
def response_volume(sat):
    """Compute first-octave responses at every valid pixel, with NaN borders."""
    sat = _validate_sat(sat)
    height, width = sat.shape[0] - 1, sat.shape[1] - 1
    responses = np.full((len(FILTER_SIZES), height, width), np.nan)
    signs = np.zeros(responses.shape, dtype=np.int8)
    for level, size in enumerate(FILTER_SIZES):
        border = size // 2
        for y in range(border, height - border):
            for x in range(border, width - border):
                *_, determinant, sign = box_hessian(sat, y, x, size)
                responses[level, y, x] = determinant
                signs[level, y, x] = sign
    return responses, signs


def scale_maxima(responses, signs, threshold=1e-5):
    """Return strict interior 3x3x3 maxima as (level, y, x, value, sign)."""
    responses, signs = np.asarray(responses), np.asarray(signs)
    if (
        responses.ndim != 3
        or responses.shape[0] != len(FILTER_SIZES)
        or signs.shape != responses.shape
        or responses.dtype.kind not in "iuf"
        or not np.isin(signs, [-1, 0, 1]).all()
        or np.isinf(responses).any()
        or not math.isfinite(threshold)
        or threshold < 0
    ):
        raise ValueError(
            "Expected four response layers, matching signs, and a finite threshold >= 0."
        )
    maxima = []
    for level in range(1, len(FILTER_SIZES) - 1):
        for y in range(1, responses.shape[1] - 1):
            for x in range(1, responses.shape[2] - 1):
                value = responses[level, y, x]
                neighbors = responses[
                    level - 1 : level + 2, y - 1 : y + 2, x - 1 : x + 2
                ].ravel()
                others = np.delete(neighbors, 13)
                if (
                    np.isfinite(neighbors).all()
                    and value > threshold
                    and np.all(value > others)
                ):
                    maxima.append(
                        (level, y, x, float(value), int(signs[level, y, x]))
                    )
    return sorted(maxima, key=lambda point: -point[3])


# endregion: scale_maxima


# region: haar
def haar_response(sat, y, x, size):
    """Return right-minus-left dx and bottom-minus-top dy for an even square."""
    sat = _validate_sat(sat)
    if (
        not _integer(size)
        or size < 2
        or size % 2
        or not _integer(y)
        or not _integer(x)
    ):
        raise ValueError(
            "Haar size must be positive and even; center must be integer."
        )
    half = size // 2
    if not (
        half <= y <= sat.shape[0] - 1 - half
        and half <= x <= sat.shape[1] - 1 - half
    ):
        raise ValueError("The full Haar wavelet must fit inside the image.")
    dx = _rectangle(sat, y - half, x, y + half, x + half)
    dx -= _rectangle(sat, y - half, x - half, y + half, x)
    dy = _rectangle(sat, y, x - half, y + half, x + half)
    dy -= _rectangle(sat, y - half, x - half, y, x + half)
    return dx, dy


# endregion: haar


# region: orientation
def orientation_from_responses(vectors):
    """Choose the longest sum in a half-open pi/3 angular window.

    Candidate starts at each nonzero response angle enumerate distinct maximal
    membership sets, including wraparound; first maximal candidate breaks ties.
    """
    vectors = _real_array(vectors)
    if (
        vectors.ndim != 2
        or vectors.shape[1] != 2
        or not np.isfinite(vectors).all()
    ):
        raise ValueError("Expected finite wavelet vectors [N, 2].")
    vectors = vectors[np.any(vectors != 0, axis=1)]
    if not len(vectors):
        return 0.0
    angles = np.mod(np.arctan2(vectors[:, 1], vectors[:, 0]), 2 * math.pi)
    best, length = np.zeros(2), -1.0
    for start in angles:
        delta = np.mod(angles - start, 2 * math.pi)
        resultant = vectors[delta < math.pi / 3].sum(axis=0)
        squared_length = float(resultant @ resultant)
        if squared_length > length:
            best, length = resultant, squared_length
    return float(math.atan2(best[1], best[0]) % (2 * math.pi))


def dominant_orientation(sat, y, x, scale):
    """Sample radius 6s at step s with 4s Haar filters and sigma 2.5s."""
    if not _integer(scale) or scale < 1:
        raise ValueError(
            "This integer-grid helper requires integer scale >= 1."
        )
    vectors = []
    for iy in range(-6, 7):
        for ix in range(-6, 7):
            if ix**2 + iy**2 <= 36:
                dx, dy = haar_response(
                    sat, y + iy * scale, x + ix * scale, 4 * scale
                )
                weight = math.exp(-(ix**2 + iy**2) / (2 * 2.5**2))
                vectors.append([weight * dx, weight * dy])
    return orientation_from_responses(vectors)


# endregion: orientation


# region: descriptor
def upright_descriptor(sat, y, x, scale):
    """Return a complete upright 64-D descriptor and its unnormalized sums.

    Integer scale s: samples use offsets [-10s, ..., 9s] in both axes. Each of
    4x4 cells has 5x5 samples; Haar size is 2s, Gaussian sigma is 3.3s. No rotation
    of coordinates or responses is applied. All wavelets must lie inside the
    image; a zero-energy patch returns a finite zero descriptor.
    """
    if not _integer(scale) or scale < 1:
        raise ValueError("The upright descriptor requires integer scale >= 1.")
    cells = np.zeros((4, 4, 4))
    for row in range(4):
        for col in range(4):
            for iy in range(5):
                for ix in range(5):
                    oy, ox = (5 * row + iy - 10) * scale, (
                        5 * col + ix - 10
                    ) * scale
                    dx, dy = haar_response(sat, y + oy, x + ox, 2 * scale)
                    weight = math.exp(
                        -(ox**2 + oy**2) / (2 * (3.3 * scale) ** 2)
                    )
                    cells[row, col] += weight * np.array(
                        [dx, dy, abs(dx), abs(dy)]
                    )
    raw = cells.ravel()
    norm = float(np.linalg.norm(raw))
    descriptor = raw / norm if norm > 0 else np.zeros(64)
    return descriptor, raw


# endregion: descriptor


# region: matching
def match_descriptors(queries, query_signs, targets, target_signs, ratio=0.8):
    """Return (query, target, ratio) after same-sign gating and strict ratio test.

    At least two eligible targets are required. Equal nearest distances and
    zero/zero ties are rejected. The threshold 0.8 is a declared example choice,
    not a guaranteed paper operating point. Distances are Euclidean, not squared.
    """
    queries, targets = _real_array(queries), _real_array(targets)
    query_signs, target_signs = np.asarray(query_signs), np.asarray(
        target_signs
    )
    for descriptors, signs in [(queries, query_signs), (targets, target_signs)]:
        if (
            descriptors.ndim != 2
            or descriptors.shape[1] != 64
            or signs.shape != (len(descriptors),)
            or not np.isfinite(descriptors).all()
            or not np.isin(signs, [-1, 0, 1]).all()
        ):
            raise ValueError(
                "Expected finite [N,64] descriptors and N trace signs."
            )
    if not math.isfinite(ratio) or not 0 < ratio < 1:
        raise ValueError(
            "Ratio threshold must lie strictly between zero and one."
        )
    matches = []
    for i, query in enumerate(queries):
        eligible = np.flatnonzero(target_signs == query_signs[i])
        if len(eligible) < 2:
            continue
        distances = np.linalg.norm(targets[eligible] - query, axis=1)
        order = np.argsort(distances, kind="stable")
        first, second = distances[order[:2]]
        if first < ratio * second:
            matches.append((i, int(eligible[order[0]]), float(first / second)))
    return matches


# endregion: matching


def synthetic_image(seed=7, size=73):
    """Generate seeded textured Gaussian blobs, without external data."""
    rng = np.random.default_rng(seed)
    y, x = np.mgrid[:size, :size]
    image = 0.01 * rng.normal(size=(size, size)) + 0.002 * x + 0.001 * y
    for cy, cx, sigma, amplitude in [(24, 25, 2.5, 1), (45, 47, 3.5, 0.8)]:
        image += amplitude * np.exp(
            -((x - cx) ** 2 + (y - cy) ** 2) / (2 * sigma**2)
        )
    return image


def smoke(seed):
    """Run measurable CPU operations and return finite JSON observations."""
    started = time.process_time()
    image = synthetic_image(seed)
    sat = integral_image(image)
    responses, signs = response_volume(sat)
    points = scale_maxima(responses, signs)
    if not points:
        raise RuntimeError("Synthetic fixture yielded no scale-space maxima.")
    # Select a detected point whose full descriptor/orientation support fits.
    eligible = []
    for point in points:
        point_scale = int(round(1.2 * FILTER_SIZES[point[0]] / 9))
        if all(
            11 * point_scale
            <= coordinate
            <= image.shape[axis] - 10 * point_scale
            for axis, coordinate in enumerate(point[1:3])
        ):
            eligible.append(point)
    if not eligible:
        raise RuntimeError(
            "No detected point has complete integer-grid support."
        )
    level, y, x, response, sign = eligible[0]
    continuous_scale = 1.2 * FILTER_SIZES[level] / 9
    scale = int(round(continuous_scale))
    descriptor, raw = upright_descriptor(sat, y, x, scale)
    changed, changed_raw = upright_descriptor(
        integral_image(3 * image + 17), y, x, scale
    )
    orientation = dominant_orientation(sat, y, x, scale)
    distractor, _ = upright_descriptor(sat, 36, 36, 1)
    matches = match_descriptors(
        descriptor[None], [sign], np.stack([changed, distractor]), [sign, sign]
    )
    return {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "seed": seed,
        "device": "cpu",
        "image_shape": list(image.shape),
        "integral_shape": list(sat.shape),
        "response_shape": list(responses.shape),
        "filter_sizes": list(FILTER_SIZES),
        "keypoint_count": len(points),
        "selected_keypoint": [level, y, x],
        "response_peak": response,
        "continuous_detector_scale": continuous_scale,
        "integer_descriptor_scale": scale,
        "descriptor_shape": list(descriptor.shape),
        "descriptor_norm": float(np.linalg.norm(descriptor)),
        "raw_descriptor_norm": float(np.linalg.norm(raw)),
        "laplacian_sign": sign,
        "orientation_radians": orientation,
        "contrast_descriptor_error": float(
            np.max(np.abs(descriptor - changed))
        ),
        "raw_contrast_error": float(np.max(np.abs(changed_raw - 3 * raw))),
        "matches": matches,
        "elapsed_cpu_seconds": time.process_time() - started,
        "reductions": [
            "73x73 synthetic image",
            "first octave only",
            "integer pixel maxima",
            "integer descriptor scale",
            "upright 64-D sampling",
            "no rotated sampling",
            "no subpixel refinement or scale interpolation",
        ],
    }


def main():
    """Print smoke observations only on explicit command-line execution."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    print(json.dumps(smoke(args.seed), allow_nan=False, sort_keys=True))


if __name__ == "__main__":
    main()