"""Fingerprint minutiae extraction v2: adaptive thresholding + skeleton.

v1 problem: global Otsu failed on low-contrast scans -> ragged blob mask ->
skeleton traced blob edges (rails/crossbars) instead of ridge centers.

v2: local adaptive threshold (window ~ ridge period x 4), ridge polarity via
global intensity stats, heavier denoise, same graph-based event detection.
"""
import sys
import cv2
import numpy as np
from skimage.morphology import skeletonize
from collections import deque

src = sys.argv[1]
out = sys.argv[2]
debug = len(sys.argv) > 3 and sys.argv[3] == 'debug'

img = cv2.imread(src, cv2.IMREAD_GRAYSCALE)
if img is None:
    print("FAIL: cannot read", src)
    sys.exit(1)
h, w = img.shape
print(f"image {w}x{h}")

# --- 1. denoise + adaptive binarization ---
den = cv2.GaussianBlur(img, (5, 5), 0)
# estimate ridge period: rough autocorrelation-free guess ~ 12-20 px for 640px scans
win = 41  # must be odd; ~2-3x ridge period
ridge = cv2.adaptiveThreshold(den, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                              cv2.THRESH_BINARY, 51, 3)
# in adaptive: 255 = darker than local mean. Which polarity is ridge?
# ridges are dark lines on light ground: dark pixels = ridge
ridge = (ridge == 0).astype(np.uint8)
frac = ridge.mean()
print(f"ridge fraction: {frac:.3f}")
# if the dark fraction is implausibly high/low for a print region, the print
# may occupy only part of the frame; restrict to the print region via
# global Otsu on the denoised image (print region is darker than bg)
_, gotsu = cv2.threshold(den, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
print_region = (gotsu == 0).astype(np.uint8)
if print_region.mean() > 0.5:
    print_region = 1 - print_region
print(f"print region fraction: {print_region.mean():.3f}")
ridge = (ridge & print_region).astype(np.uint8)

# --- 2. cleanup ---
k2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))
k3 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
k5 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
ridge = cv2.morphologyEx(ridge, cv2.MORPH_OPEN, k2, iterations=1)   # kill grain
ridge = cv2.morphologyEx(ridge, cv2.MORPH_CLOSE, k3, iterations=1)  # bridge grain gaps

# keep components with enough area
n_lab, labels, stats, _ = cv2.connectedComponentsWithStats(ridge, connectivity=8)
keep = np.zeros_like(ridge)
for i in range(1, n_lab):
    if stats[i, cv2.CC_STAT_AREA] >= 100:
        keep[labels == i] = 1
ridge = keep
print(f"ridge fraction after cleanup: {ridge.mean():.3f}")

if debug:
    canvas = np.hstack([img, (ridge * 255).astype(np.uint8)])
    cv2.imwrite('/tmp/fp_debug2.png', canvas)

# --- 3. skeletonize ---
skel = skeletonize(ridge > 0).astype(np.uint8)
print(f"skeleton pixels: {skel.sum()}")

# prune short spurs (< 6 px) from skeleton
ys, xs = np.nonzero(skel)
coords = set(zip(ys.tolist(), xs.tolist()))
adj = {}
for (y, x) in coords:
    adj[(y, x)] = [(y + dy, x + dx) for dy in (-1, 0, 1) for dx in (-1, 0, 1)
                   if (dy, dx) != (0, 0) and (y + dy, x + dx) in coords]

def prune_spurs(skel_pts, adj, min_len=6, iters=3):
    pts = set(skel_pts)
    for _ in range(iters):
        removed = set()
        # find endpoints
        ends = [p for p in pts if len(adj[p]) == 1]
        for e in ends:
            # walk branch until branch point or end
            path = [e]
            cur = e
            prev = None
            while True:
                nxts = [n for n in adj[cur] if n != prev]
                if len(nxts) != 1:
                    break
                nxt = nxts[0]
                if nxt in pts:
                    path.append(nxt)
                    prev, cur = cur, nxt
                else:
                    break
            if len(path) < min_len:
                removed.update(path)
        if not removed:
            break
        for p in removed:
            pts.discard(p)
        # rebuild adj for remaining
        for p in pts:
            adj[p] = [n for n in adj[p] if n in pts]
    return pts

coords = prune_spurs(coords, adj, min_len=6, iters=2)
skel = np.zeros_like(ridge)
for (y, x) in coords:
    skel[y, x] = 1
print(f"skeleton pixels after spur pruning: {skel.sum()}")

# --- 4. candidate events ---
cands = []
for (y, x) in coords:
    n = len(adj[(y, x)])
    if n == 1:
        cands.append((y, x, 'end'))
    elif n >= 3:
        cands.append((y, x, 'bif'))
print(f"raw candidates: end={sum(1 for c in cands if c[2]=='end')}, bif={sum(1 for c in cands if c[2]=='bif')}")

# --- 4b. branch-length filter ---
event_set = {(c[0], c[1]) for c in cands}

def branch_lengths(pt):
    lengths = []
    for nb in adj[pt]:
        if nb in event_set:
            continue
        visited = {pt, nb}
        q = deque([(nb, 0)])
        maxlen = 0
        while q:
            (cy, cx), d = q.popleft()
            maxlen = max(maxlen, d)
            for nn in adj[(cy, cx)]:
                if nn in visited:
                    continue
                if nn in event_set:
                    maxlen = max(maxlen, d + 1)
                    continue
                visited.add(nn)
                q.append((nn, d + 1))
        lengths.append(maxlen)
    return lengths

MIN_BRANCH = 12
filtered = []
for (y, x, t) in cands:
    bl = branch_lengths((y, x))
    if not bl:
        filtered.append((y, x, t))
        continue
    if t == 'bif':
        # require at least 3 branches of length >= MIN_BRANCH (a real Y,
        # not a grain knot)
        if sum(1 for L in bl if L >= MIN_BRANCH) < 3:
            continue
    else:
        if len(bl) == 1 and bl[0] < MIN_BRANCH:
            continue
    filtered.append((y, x, t))
print(f"after branch-length filter: end={sum(1 for c in filtered if c[2]=='end')}, bif={sum(1 for c in filtered if c[2]=='bif')}")
cands = filtered

# --- 5. cluster nearby candidates ---
def cluster(cands, radius=3):
    used = [False] * len(cands)
    groups = []
    for i, c in enumerate(cands):
        if used[i]:
            continue
        stack = [i]
        used[i] = True
        group = [c]
        while stack:
            j = stack.pop()
            yj, xj, _ = cands[j]
            for k in range(len(cands)):
                if not used[k] and abs(cands[k][0] - yj) <= radius and abs(cands[k][1] - xj) <= radius:
                    used[k] = True
                    stack.append(k)
                    group.append(cands[k])
        groups.append(group)
    return groups

events = []
for g in cluster(cands):
    y = sum(c[0] for c in g) / len(g)
    x = sum(c[1] for c in g) / len(g)
    t = 'bif' if any(c[2] == 'bif' for c in g) else 'end'
    # anchor = first member (a real skeleton pixel with adjacency info)
    events.append((int(y), int(x), t, g[0][0], g[0][1]))
print(f"after clustering: end={sum(1 for e in events if e[2]=='end')}, bif={sum(1 for e in events if e[2]=='bif')}")

# --- 6. spatial filters ---
dist = cv2.distanceTransform(ridge, cv2.DIST_L2, 5)
def ridge_width(y, x):
    y0, y1 = max(0, y - 6), min(h, y + 7)
    x0, x1 = max(0, x - 6), min(w, x + 7)
    return 2 * dist[y0:y1, x0:x1].max()

blob = cv2.dilate(ridge, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15)), iterations=2)
dist_to_outer = cv2.distanceTransform(blob, cv2.DIST_L2, 5)

final = []
for (y, x, t, ay, ax) in events:
    if not ridge[y, x]:
        continue  # safety: must be on a ridge
    if ridge_width(y, x) < 3.0 or ridge_width(y, x) > 12.0:
        continue
    # boundary ring: events near the outer edge of the print are unreliable
    if dist_to_outer[y, x] < (22 if t == 'end' else 14):
        continue
    if t == 'end':
        # break-detection: if ridge pixels continue nearby in the forward
        # direction, this "end" is a grain gap, not a true ending
        if (ay, ax) not in adj or len(adj[(ay, ax)]) != 1:
            continue
        # skeleton direction at the endpoint = away from its only neighbor
        nb = adj[(ay, ax)][0]
        dy, dx = ay - nb[0], ax - nb[1]
        n = (dy * dy + dx * dx) ** 0.5
        if n < 1e-6:
            continue
        dy, dx = dy / n, dx / n
        # grain-gap test on the ridge MASK: if the ridge resumes within a
        # corridor (±2px) 6..20 px ahead, the break is grain.
        hit = False
        for s in range(6, 21):
            cy, cx = ay + dy * s, ax + dx * s
            for off in (-2, -1, 0, 1, 2):
                fy = int(round(cy - dx * off))
                fx = int(round(cx + dy * off))
                if 0 <= fy < h and 0 <= fx < w and ridge[fy, fx]:
                    hit = True
                    break
            if hit:
                break
        if hit:
            continue
    final.append((x, y, t))
print(f"after spatial filters: end={sum(1 for e in final if e[2]=='end')}, bif={sum(1 for e in final if e[2]=='bif')}")

# --- 7. density cap ---
cell = 20
grid = {}
capped = []
for (x, y, t) in sorted(final, key=lambda e: (e[2] != 'bif', e[0], e[1])):
    key = (x // cell, y // cell)
    if key in grid:
        continue
    grid[key] = t
    capped.append((x, y, t))
print(f"after density cap: end={sum(1 for e in capped if e[2]=='end')}, bif={sum(1 for e in capped if e[2]=='bif')}")

# --- 8. draw ---
vis = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
for i, (x, y, t) in enumerate(capped, 1):
    if t == 'end':
        cv2.circle(vis, (x, y), 9, (0, 255, 0), 2)
        cv2.putText(vis, str(i), (x + 10, y - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 200, 0), 1)
    else:
        cv2.rectangle(vis, (x - 8, y - 8), (x + 8, y + 8), (0, 0, 255), 2)
        cv2.putText(vis, str(i), (x + 10, y - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 220), 1)
# legend (bottom-right, on the blank margin)
lx, ly = w - 235, h - 95
cv2.rectangle(vis, (lx, ly), (w - 12, h - 12), (255, 255, 255), -1)
cv2.rectangle(vis, (lx, ly), (w - 12, h - 12), (0, 0, 0), 1)
cv2.circle(vis, (lx + 22, ly + 25), 9, (0, 255, 0), 2)
cv2.putText(vis, f"endpoint  {sum(1 for e in capped if e[2]=='end')}", (lx + 42, ly + 30),
            cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 1)
cv2.rectangle(vis, (lx + 14, ly + 52), (lx + 30, ly + 68), (0, 0, 255), 2)
cv2.putText(vis, f"bifurcation  {sum(1 for e in capped if e[2]=='bif')}", (lx + 42, ly + 63),
            cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 1)
cv2.imwrite(out, vis, [cv2.IMWRITE_PNG_COMPRESSION, 3])
# also dump coordinates (CSV) next to the output
import csv as _csv
with open(out + '.csv', 'w', newline='') as f:
    wr = _csv.writer(f)
    wr.writerow(['id', 'x', 'y', 'type'])
    for i, (x, y, t) in enumerate(capped, 1):
        wr.writerow([i, x, y, t])
print(f"WROTE {out}")
print(f"TOTAL: {len(capped)} minutiae (end={sum(1 for e in capped if e[2]=='end')}, bif={sum(1 for e in capped if e[2]=='bif')})")