import fitz, math
from collections import defaultdict

TOL = 0.6
WIRE_COLORS = {(1.0,0.0,0.0), (0.5,0.0,0.25)}   # red pin stubs + maroon wires

class DSU:
    def __init__(s): s.p={}
    def find(s,x):
        s.p.setdefault(x,x)
        while s.p[x]!=x:
            s.p[x]=s.p[s.p[x]]; x=s.p[x]
        return x
    def union(s,a,b):
        ra,rb=s.find(a),s.find(b)
        if ra!=rb: s.p[ra]=rb

def segs_of_page(page):
    out=[]
    for d in page.get_drawings():
        col=d.get('color')
        if col is None: continue
        c=tuple(round(v,3) for v in col)
        if c not in WIRE_COLORS: continue
        for it in d['items']:
            if it[0]=='l':
                p,q=it[1],it[2]
                out.append(((p.x,p.y),(q.x,q.y)))
    return [s for s in out if math.dist(s[0],s[1])>0.05]

def on_seg(p,a,b,tol=TOL):
    (px,py),(ax,ay),(bx,by)=p,a,b
    dx,dy=bx-ax,by-ay
    L2=dx*dx+dy*dy
    if L2==0: return False
    t=((px-ax)*dx+(py-ay)*dy)/L2
    if t<-1e-9 or t>1+1e-9: return False
    cx,cy=ax+t*dx,ay+t*dy
    return math.hypot(px-cx,py-cy)<=tol


def junction_dots(page):
    out=[]
    for d in page.get_drawings():
        f=d.get('fill')
        if d['type']=='fs' and f and abs(f[0]-1)<.01 and f[1]<.01 and abs(f[2]-1)<.01:
            r=d['rect']; out.append(((r.x0+r.x1)/2,(r.y0+r.y1)/2))
    return out

def nets(page):
    segs=segs_of_page(page)
    dsu=DSU()
    key=lambda p:(round(p[0]*2)/2, round(p[1]*2)/2)
    for a,b in segs: dsu.union(key(a),key(b))
    for i,(a,b) in enumerate(segs):
        for p in (a,b):
            for j,(c,e) in enumerate(segs):
                if i==j: continue
                if on_seg(p,c,e):
                    dsu.union(key(p),key(c)); dsu.union(key(p),key(e))
    for dc in junction_dots(page):
        hit=[]
        for a,b in segs:
            if on_seg(dc,a,b,tol=1.4): hit.append((a,b))
        for a,b in hit[1:]:
            dsu.union(key(a),key(hit[0][0])); dsu.union(key(b),key(hit[0][0]))
    comp=defaultdict(set)
    for a,b in segs:
        for p in (a,b): comp[dsu.find(key(p))].add(key(p))
    return segs,dsu,key,comp