"""Run: python verify_cv.py. No downloads, GPU, or training required."""
from __future__ import annotations

from collections.abc import Callable
import json
import platform

import cv2
import numpy as np
import scipy
from scipy.ndimage import distance_transform_edt, label as scipy_label
from scipy.spatial.distance import directed_hausdorff

from cv_channels import FloatImage, channel_planes, channel_statistics, alpha_example, composite_over, srgb_to_linear, linear_to_srgb
from cv_components import Connectivity, component_scene, label_components, measure_component, keep_by_area
from cv_filters import FloatPlane, inspect_patch, smoothing_kernel, step_image, correlate3, smooth_opencv, sobel_components
from cv_arrays import pixel_demo, to_unit_rgb, model_input
from cv_masks import scene, threshold_mask, clean_mask
from cv_metrics import Mask, boundary, directed_distances, hausdorff, metric_cases, overlap


def rejects(run: Callable[[], object]) -> None:
    try:
        run()
    except ValueError:
        return
    raise AssertionError("Expected input contract rejection")


def slow_boundary(mask: Mask) -> Mask:
    """Independent Python neighbor check, with outside-image background."""
    result: Mask = np.zeros_like(mask)
    height, width = mask.shape
    for row in range(height):
        for col in range(width):
            if not mask[row, col]:
                continue
            for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                rr, cc = row + dr, col + dc
                if not (0 <= rr < height and 0 <= cc < width) or not mask[rr, cc]:
                    result[row, col] = True
                    break
    return result


def check_components() -> None:
    mask = component_scene()
    assert int(mask.sum()) == 13
    for connectivity in (4, 8):
        mode: Connectivity = 4 if connectivity == 4 else 8
        count, labels = label_components(mask,mode)
        assert count == (5 if mode == 4 else 3)
        kept, areas = keep_by_area(labels,3)
        assert sorted(areas[1:].tolist()) == ([1,1,1,4,6] if mode == 4 else [1,6,6])
        assert int(kept.sum()) == (10 if mode == 4 else 12)
        assert not kept[~mask].any()
        assert np.array_equal(keep_by_area(labels,1)[0],mask)
        for component_id in range(1,count+1):
            area, center, box = measure_component(labels,component_id)
            coords = [(r,c) for r in range(7) for c in range(9) if labels[r,c] == component_id]
            assert area == len(coords)
            np.testing.assert_allclose(center,[sum(p[0] for p in coords)/area,sum(p[1] for p in coords)/area])
            assert box == (min(p[0] for p in coords),min(p[1] for p in coords),max(p[0] for p in coords)+1,max(p[1] for p in coords)+1)
    # Independent SciPy partition oracle; do not assume shared label numbering.
    rng = np.random.default_rng(43)
    for mode in (4,8):
        structure = np.ones((3,3),dtype=np.uint8) if mode == 8 else np.array([[0,1,0],[1,1,1],[0,1,0]],dtype=np.uint8)
        for _ in range(20):
            random_mask = rng.random((6,9)) > .65
            actual_count, actual = label_components(random_mask,mode)
            reference, reference_count = scipy_label(random_mask,structure=structure)
            assert actual_count == reference_count
            np.testing.assert_array_equal(actual > 0,reference > 0)
            # Each partition is the same even when the assigned IDs differ.
            mapping: dict[int,int] = {}
            for component_id in range(1,actual_count+1):
                reference_ids = np.unique(reference[actual == component_id])
                assert reference_ids.size == 1 and reference_ids[0] != 0
                mapping[component_id] = int(reference_ids[0])
            assert len(set(mapping.values())) == actual_count
    empty = np.zeros((2,3),dtype=np.bool_)
    count, labels = label_components(empty)
    assert count == 0 and not keep_by_area(labels,1)[0].any()
    full = np.ones((2,3),dtype=np.bool_)
    count, labels = label_components(full)
    assert count == 1 and keep_by_area(labels,6)[0].all()
    assert not keep_by_area(labels,7)[0].any()
    assert measure_component(labels,1) == (6,(.5,1.0),(0,0,2,3))
    rejects(lambda: measure_component(labels,0))
    rejects(lambda: measure_component(labels,2))
    rejects(lambda: keep_by_area(labels,0))
    rejects(lambda: keep_by_area(np.full((2,3),99,dtype=np.int32),1))
    rejects(lambda: label_components(full.astype(np.uint8)))
    print("PASS: connected-component counts, area selection, centroids and boxes; 40 independent SciPy partition checks; empty/full masks and invalid contracts")


def check_filters() -> None:
    kernel = smoothing_kernel()
    patch = np.array([[0,0,80]] * 3, dtype=np.float32)
    products, response = inspect_patch(patch, kernel)
    np.testing.assert_array_equal(products, [[0,0,5],[0,0,10],[0,0,5]])
    assert response == 20.0
    step = step_image()
    np.testing.assert_array_equal(correlate3(step,kernel), [[0,20,60,80,80]] * 3)
    gx, gy, magnitude = sobel_components(step)
    np.testing.assert_array_equal(gx, [[0,320,320,0,0]] * 3)
    np.testing.assert_array_equal(gy, np.zeros_like(step))
    np.testing.assert_array_equal(magnitude, gx)
    falling = sobel_components(80-step)
    np.testing.assert_array_equal(falling[0], -gx)
    np.testing.assert_array_equal(falling[2], magnitude)
    ramp: FloatPlane = np.tile(np.arange(7,dtype=np.float32), (5,1))
    np.testing.assert_array_equal(sobel_components(ramp)[0][:,1:-1], 8)
    constant = np.full((4,7),50,dtype=np.float32)
    np.testing.assert_allclose(correlate3(constant,kernel),constant)
    assert not sobel_components(constant)[2].any()
    # Independent direct scalar accumulation, asymmetric kernel detects flips.
    rng = np.random.default_rng(37)
    for shape in [(1,1),(2,7),(6,5)]:
        image: FloatPlane = rng.normal(size=shape).astype(np.float32)
        asymmetric: FloatPlane = rng.normal(size=(3,3)).astype(np.float32)
        oracle = np.zeros_like(image)
        for row in range(shape[0]):
            for col in range(shape[1]):
                value: float = 0.0
                for kr in range(3):
                    for kc in range(3):
                        rr = min(max(row+kr-1,0),shape[0]-1)
                        cc = min(max(col+kc-1,0),shape[1]-1)
                        value += float(image[rr,cc]) * float(asymmetric[kr,kc])
                oracle[row,col] = value
        np.testing.assert_allclose(correlate3(image,asymmetric),oracle,atol=2e-6)
        np.testing.assert_allclose(smooth_opencv(image),correlate3(image,kernel),atol=2e-7)
        direct = cv2.filter2D(image,cv2.CV_32F,asymmetric,borderType=cv2.BORDER_REPLICATE)
        np.testing.assert_allclose(direct,oracle,atol=2e-6)
    rejects(lambda: correlate3(constant,np.ones((2,2),dtype=np.float32)))
    rejects(lambda: sobel_components(constant.astype(np.uint8)))
    rejects(lambda: sobel_components(np.full((2,2),np.nan,dtype=np.float32)))
    print("PASS: filtering worked example; replicated borders; signed step and ramp gradients; independent scalar correlation oracle; OpenCV agreement; invalid contracts")


def check_channels() -> None:
    bgr, rgb, _ = pixel_demo()
    red, green, blue, rebuilt = channel_planes(rgb)
    np.testing.assert_array_equal(rebuilt, rgb)
    np.testing.assert_array_equal(red, [[255,0],[0,90]])
    red[0,0] = 0
    assert rgb[0,0,0] == 255
    unit = to_unit_rgb(bgr)
    averages, pixels, mask, selected = channel_statistics(unit)
    np.testing.assert_allclose(averages, np.array([345,315,285]) / 1020, rtol=1e-6)
    np.testing.assert_allclose(pixels, np.array([[1/3,1/3],[1/3,60/255]], dtype=np.float32), rtol=1e-6)
    np.testing.assert_array_equal(mask, [[True,False],[False,True]])
    expected = unit.copy(); expected[0,1] = 0; expected[1,0] = 0
    np.testing.assert_array_equal(selected, expected)
    f, b, a = alpha_example()
    actual = composite_over(f,b,a)
    np.testing.assert_allclose(actual, np.array([[[0,0,1],[.25,0,.75],[.5,0,.5],[1,0,0]]], dtype=np.float32))
    # Independent scalar oracle on non-square random inputs.
    rng = np.random.default_rng(21)
    fg: FloatImage = rng.random((3,5,3), dtype=np.float32)
    bg: FloatImage = rng.random((3,5,3), dtype=np.float32)
    opacity: FloatImage = rng.random((3,5,1), dtype=np.float32)
    oracle: FloatImage = np.zeros_like(fg)
    for row in range(3):
        for col in range(5):
            for channel in range(3):
                weight = float(opacity[row,col,0])
                oracle[row,col,channel] = weight * float(fg[row,col,channel]) + (1-weight) * float(bg[row,col,channel])
    np.testing.assert_allclose(composite_over(fg,bg,opacity), oracle, atol=1e-7)
    np.testing.assert_array_equal(composite_over(fg,bg,np.zeros_like(opacity)), bg)
    np.testing.assert_array_equal(composite_over(fg,bg,np.ones_like(opacity)), fg)
    rejects(lambda: composite_over(fg,bg,opacity[:1]))
    rejects(lambda: composite_over(fg,bg,-np.ones_like(opacity)))
    rejects(lambda: composite_over(fg,bg,np.full_like(opacity,np.nan)))
    rejects(lambda: channel_planes(rgb.astype(np.float32)))
    np.testing.assert_allclose(srgb_to_linear(linear_to_srgb(fg)),fg,atol=3e-7)
    np.testing.assert_allclose(linear_to_srgb(srgb_to_linear(fg)),fg,atol=3e-7)
    np.testing.assert_allclose(linear_to_srgb(np.full((1,1,3),.5,dtype=np.float32)),.73535698,atol=1e-7)
    np.testing.assert_allclose(srgb_to_linear(np.full((1,1,3),.04045,dtype=np.float32)),.04045/12.92,atol=1e-8)
    print("PASS: channel planes, axis meanings, selection, alpha endpoints, scalar compositing oracle, color transfer and invalid contracts")


def check() -> None:
    check_channels()
    check_filters()
    check_components()
    bgr, rgb, crop = pixel_demo()
    assert rgb[0, 0].tolist() == [255, 0, 0]
    assert rgb[1, 1].tolist() == [90, 60, 30] and crop[0, 0].tolist() == [0, 60, 30]
    assert not np.shares_memory(crop, rgb)
    batch = model_input(to_unit_rgb(bgr))
    assert batch.shape == (1, 3, 2, 2) and batch.flags.c_contiguous
    np.testing.assert_allclose(batch[0, :, 1, 1], [90/255, 60/255, 30/255])
    rejects(lambda: to_unit_rgb(bgr.astype(np.float32)))
    rejects(lambda: model_input(rgb.astype(np.float32)))
    gray, target = scene()
    raw = threshold_mask(gray)
    closed, cleaned = clean_mask(raw)
    assert raw.sum() == 36 and target.sum() == 36
    assert not raw[5, 5] and raw[1, 1]
    np.testing.assert_array_equal(cleaned, target)
    edge_values = np.array([[99, 100, 101]], dtype=np.uint8)
    np.testing.assert_array_equal(threshold_mask(edge_values), [[False, False, True]])
    # Cleaning is not universally safe: a legitimate single-pixel object vanishes.
    thin = np.zeros((7, 7), dtype=np.bool_)
    thin[3, 3] = True
    assert not clean_mask(thin)[1].any()
    results: dict[str, dict[str, float | str]] = {}
    for name, (pred, truth) in metric_cases().items():
        iou, dice = overlap(pred, truth)
        hd, hd95 = hausdorff(pred, truth)
        np.testing.assert_allclose(dice, 2 * iou / (1 + iou))
        np.testing.assert_allclose(overlap(pred, truth), overlap(truth, pred))
        np.testing.assert_allclose(hausdorff(pred, truth), hausdorff(truth, pred))
        if pred.any() and truth.any():
            ap, at = np.argwhere(boundary(pred)), np.argwhere(boundary(truth))
            oracle = max(directed_hausdorff(ap, at)[0], directed_hausdorff(at, ap)[0])
            np.testing.assert_allclose(hd, oracle)
        results[name] = dict(zip(("iou", "dice", "hd", "hd95"),
                                [float(v) if np.isfinite(v) else "infinity" for v in (iou, dice, hd, hd95)]))
    np.testing.assert_allclose([float(value) for value in results["shifted"].values()], [5/7, 5/6, 1, 1])
    np.testing.assert_allclose([float(value) for value in results["outlier"].values()], [36/37, 72/73, np.sqrt(18), 0])
    assert results["both-empty"] == dict(iou=1.0, dice=1.0, hd=0.0, hd95=0.0)
    assert results["one-empty"] == dict(iou=0.0, dice=0.0, hd="infinity", hd95="infinity")
    # Spacing follows (row,column): a one-column shift is 3 units, not 2.
    p = np.zeros((5, 5), dtype=np.bool_); t = p.copy()
    p[2, 2] = True; t[2, 3] = True
    np.testing.assert_allclose(hausdorff(p, t, (2.0, 3.0)), [3, 3])
    rejects(lambda: hausdorff(p, t, (0.0, 1.0)))
    rejects(lambda: overlap(p, t[:2]))
    rejects(lambda: overlap(p.astype(np.uint8), t))
    # Independent EDT oracle on random masks, anisotropic and unit spacing.
    rng = np.random.default_rng(17)
    for _ in range(24):
        p, t = rng.random((2, 9, 11)) > 0.65
        np.testing.assert_array_equal(boundary(p), slow_boundary(p))
        for spacing in [(1.0, 1.0), (2.0, 0.7)]:
            ref = distance_transform_edt(~slow_boundary(t), sampling=spacing)[slow_boundary(p)]
            np.testing.assert_allclose(directed_distances(p, t, spacing), ref)
    full = np.ones((4, 4), dtype=np.bool_)
    np.testing.assert_array_equal(boundary(full), slow_boundary(full))
    print(json.dumps({"environment": {"python": platform.python_version(), "numpy": np.__version__,
          "opencv": cv2.__version__, "scipy": scipy.__version__}, "cases": results}, indent=2))
    print("PASS: image/layout contracts; threshold equality; cleaning success/failure; six metric cases; symmetry; spacing; SciPy HD and 48 EDT comparisons")


if __name__ == "__main__":
    check()