#!/usr/bin/env python3
"""
vectorize_skeleton.py
Trace skeleton PNG to Manhattan polylines (simple 4-neighbor tracing) and export JSON.
"""
import argparse
from pathlib import Path
import numpy as np
from PIL import Image
import json
import cv2
from collections import deque


def load_skeleton(path):
    im = Image.open(path).convert('L')
    a = np.array(im)
    return (a > 127).astype('uint8')


def neighbors4(y, x, h, w):
    for dy, dx in ((0,1),(1,0),(0,-1),(-1,0)):
        ny, nx = y+dy, x+dx
        if 0 <= ny < h and 0 <= nx < w:
            yield ny, nx


def trace_components(sk):
    h, w = sk.shape
    visited = np.zeros_like(sk, dtype=bool)
    comps = []
    for y in range(h):
        for x in range(w):
            if sk[y,x] and not visited[y,x]:
                # BFS to collect component pixels
                q = deque()
                q.append((y,x))
                visited[y,x] = True
                pts = []
                while q:
                    cy, cx = q.popleft()
                    pts.append((int(cx), int(cy)))
                    for ny, nx in neighbors4(cy, cx, h, w):
                        if sk[ny,nx] and not visited[ny,nx]:
                            visited[ny,nx] = True
                            q.append((ny,nx))
                comps.append(pts)
    return comps


def prune_segments(segments, min_len=5, merge_tol=3):
    if not segments: return []
    
    # 1. Merge consecutive same-direction segments (simple pass)
    # Repeat until no more merges to handle multi-segment merges
    while True:
        merged = []
        changed = False
        if segments:
            curr = segments[0]
            for i in range(1, len(segments)):
                next_seg = segments[i]
                if curr[0] == next_seg[0] and abs(curr[1] - next_seg[1]) <= merge_tol:
                    # Merge: extend end_idx, re-calculate average val
                    len1 = curr[3] - curr[2] + 1
                    len2 = next_seg[3] - next_seg[2] + 1
                    new_val = int(round((curr[1]*len1 + next_seg[1]*len2) / (len1+len2)))
                    curr = (curr[0], new_val, curr[2], next_seg[3])
                    changed = True
                else:
                    merged.append(curr)
                    curr = next_seg
            merged.append(curr)
        segments = merged
        if not changed:
            break

    # 2. Remove short tails
    # Check start
    if len(segments) > 1:
        s0 = segments[0]
        l0 = s0[3] - s0[2] + 1
        if l0 < min_len:
            segments.pop(0)
    
    # Check end
    if len(segments) > 1:
        s_last = segments[-1]
        l_last = s_last[3] - s_last[2] + 1
        if l_last < min_len:
            segments.pop(-1)
            
    # 3. Remove short internal bumps (Z-shape)
    # H1 -> V(short) -> H2  => Merge H1, H2 if aligned
    # We iterate and build a new list. If merge happens, we modify the next segment and skip.
    final_segs = []
    i = 0
    while i < len(segments):
        curr = segments[i]
        merged_bump = False
        if i + 2 < len(segments):
            mid = segments[i+1]
            next_seg = segments[i+2]
            
            # Check pattern: A -> B(short) -> A
            if curr[0] == next_seg[0] and mid[0] != curr[0]:
                len_mid = mid[3] - mid[2] + 1
                if len_mid < min_len:
                    # Check alignment of curr and next_seg
                    if abs(curr[1] - next_seg[1]) <= merge_tol:
                        # Merge all three: curr + mid + next_seg -> new_curr
                        len1 = curr[3] - curr[2] + 1
                        len3 = next_seg[3] - next_seg[2] + 1
                        new_val = int(round((curr[1]*len1 + next_seg[1]*len3) / (len1+len3)))
                        # New segment spans from curr.start to next_seg.end
                        new_seg = (curr[0], new_val, curr[2], next_seg[3])
                        
                        # We effectively skip mid and next_seg, and replace curr with new_seg
                        # But we need to check if this new_seg can merge further?
                        # For simplicity, let's push new_seg to final_segs and skip 2.
                        # But wait, if we push to final_segs, we can't merge it with subsequent ones in this loop easily.
                        # Let's update segments[i+2] to be the merged one and continue loop from i+2.
                        segments[i+2] = new_seg
                        i += 2
                        merged_bump = True
                        continue
        
        if not merged_bump:
            final_segs.append(curr)
            i += 1
            
    return final_segs


def simplify_manhattan(pts, tolerance=2):
    if not pts: return []
    if len(pts) < 2: return pts
    
    segments = [] 
    n = len(pts)
    i = 0
    while i < n - 1:
        start_pt = pts[i]
        
        # Check Horizontal
        k_h = i + 1
        while k_h < n:
            if abs(pts[k_h][1] - start_pt[1]) > tolerance:
                break
            k_h += 1
        len_h = k_h - i
        
        # Check Vertical
        k_v = i + 1
        while k_v < n:
            if abs(pts[k_v][0] - start_pt[0]) > tolerance:
                break
            k_v += 1
        len_v = k_v - i
        
        if len_h >= len_v:
            # Horizontal
            segment_pts = pts[i:k_h]
            avg_y = int(round(np.mean([p[1] for p in segment_pts])))
            segments.append(('H', avg_y, i, k_h - 1))
            i = k_h - 1
        else:
            # Vertical
            segment_pts = pts[i:k_v]
            avg_x = int(round(np.mean([p[0] for p in segment_pts])))
            segments.append(('V', avg_x, i, k_v - 1))
            i = k_v - 1
    
    # --- Pruning / Refining ---
    segments = prune_segments(segments, min_len=5, merge_tol=3)
    
    if not segments:
        return []

    out_poly = []
    
    # First point
    first_seg = segments[0]
    first_pt_orig = pts[first_seg[2]]
    if first_seg[0] == 'H':
        curr = (first_pt_orig[0], first_seg[1])
    else:
        curr = (first_seg[1], first_pt_orig[1])
    out_poly.append(curr)
    
    for idx in range(len(segments) - 1):
        s1 = segments[idx]
        s2 = segments[idx+1]
        
        # Transition point index is s1[3] (which is same as s2[2])
        # But after pruning, s1[3] might not be s2[2] - 1. Gaps might exist.
        # We should just connect s1's end to s2's start via a Manhattan corner.
        
        # s1 end point (projected)
        if s1[0] == 'H':
            p1_end = (pts[s1[3]][0], s1[1])
        else:
            p1_end = (s1[1], pts[s1[3]][1])
            
        # s2 start point (projected)
        if s2[0] == 'H':
            p2_start = (pts[s2[2]][0], s2[1])
        else:
            p2_start = (s2[1], pts[s2[2]][1])
            
        # Connect p1_end to p2_start
        if s1[0] == 'H' and s2[0] == 'V':
            # H(y1) -> V(x2)
            # Intersection is (x2, y1)
            corner = (s2[1], s1[1])
            if out_poly[-1] != corner: out_poly.append(corner)
        elif s1[0] == 'V' and s2[0] == 'H':
            # V(x1) -> H(y2)
            # Intersection is (x1, y2)
            corner = (s1[1], s2[1])
            if out_poly[-1] != corner: out_poly.append(corner)
        else:
            # Parallel segments (should have been merged, but if gap was large...)
            # Or H -> H (gap)
            # Just connect via midpoint or direct L-shape?
            # Let's use p1_end -> p2_start directly? No, need Manhattan.
            # H(y1) ... H(y2). Connect (x_end1, y1) -> (x_end1, y2) -> (x_start2, y2) ?
            # Or (x_end1, y1) -> (x_start2, y1) -> (x_start2, y2) ?
            # Let's use the first one (Vertical bridge).
            if out_poly[-1] != p1_end: out_poly.append(p1_end)
            
            if s1[0] == 'H':
                # Bridge is Vertical
                mid_x = (p1_end[0] + p2_start[0]) // 2
                c1 = (mid_x, p1_end[1])
                c2 = (mid_x, p2_start[1])
                if out_poly[-1] != c1: out_poly.append(c1)
                if c1 != c2: out_poly.append(c2)
                if c2 != p2_start: out_poly.append(p2_start)
            else:
                # Bridge is Horizontal
                mid_y = (p1_end[1] + p2_start[1]) // 2
                c1 = (p1_end[0], mid_y)
                c2 = (p2_start[0], mid_y)
                if out_poly[-1] != c1: out_poly.append(c1)
                if c1 != c2: out_poly.append(c2)
                if c2 != p2_start: out_poly.append(p2_start)

    # Last point
    last_seg = segments[-1]
    last_pt_orig = pts[last_seg[3]]
    if last_seg[0] == 'H':
        end = (last_pt_orig[0], last_seg[1])
    else:
        end = (last_seg[1], last_pt_orig[1])
    if out_poly[-1] != end: out_poly.append(end)
    
    return out_poly


def polyline_from_pixels(pts):
    # pts: list of (x,y) unordered; produce simple ordering by greedy nearest neighbor
    if not pts:
        return []
    pts_set = set(pts)
    # start from leftmost-topmost
    cur = min(pts, key=lambda p: (p[1], p[0]))
    seq = [cur]
    pts_set.remove(cur)
    while pts_set:
        # look for 4-neighbor next
        x,y = seq[-1]
        found = None
        for nx, ny in ((x+1,y),(x-1,y),(x,y+1),(x,y-1)):
            if (nx,ny) in pts_set:
                found = (nx,ny)
                break
        if found is None:
            # fallback: nearest
            found = min(pts_set, key=lambda p: (abs(p[0]-x)+abs(p[1]-y), p[1], p[0]))
        seq.append(found)
        pts_set.remove(found)
    
    return simplify_manhattan(seq, tolerance=2)


def main():
    p = argparse.ArgumentParser()
    p.add_argument('skel_png')
    p.add_argument('outjson')
    args = p.parse_args()

    sk = load_skeleton(args.skel_png)
    comps = trace_components(sk)
    polylines = []
    for pts in comps:
        pl = polyline_from_pixels(pts)
        if len(pl) > 1:
            polylines.append(pl)

    out = {'polylines': polylines, 'num': len(polylines), 'shape': sk.shape}
    with open(args.outjson, 'w') as f:
        json.dump(out, f)
    print('Wrote', args.outjson)


if __name__ == '__main__':
    main()