import fitz, math
from collections import defaultdict
from trace import nets

doc=fitz.open('Xiaomi_Mi_8_Schematic.pdf')

def red_spans(page):
    out=[]
    for b in page.get_text('dict')['blocks']:
        for l in b.get('lines',[]):
            for s in l['spans']:
                if s['color']==0xFF0000:
                    t=s['text'].strip()
                    if t: out.append((t, fitz.Rect(s['bbox'])))
    return out

def gnd_components(segs,dsu,key):
    """ground symbol = >=3 short parallel segs of decreasing length sharing a component"""
    byc=defaultdict(list)
    for a,b in segs:
        byc[dsu.find(key(a))].append((a,b))
    g=set()
    for c,ss in byc.items():
        shorts=[s for s in ss if 1.0<math.dist(s[0],s[1])<9.0]
        if len(shorts)>=3:
            lens=sorted(round(math.dist(s[0],s[1]),1) for s in shorts)
            if len(set(lens))>=3: g.add(c)
    return g

def page_netmap(pgno):
    page=doc[pgno-1]
    segs,dsu,key,comp=nets(page)
    pts=[]
    for a,b in segs: pts.append(a); pts.append(b)
    lab=defaultdict(set)
    for t,r in red_spans(page):
        cands=[]
        for p in pts:
            dx=max(r.x0-p[0], 0, p[0]-r.x1); dy=max(r.y0-p[1],0,p[1]-r.y1)
            cands.append((math.hypot(dx,dy),p))
        d,p=min(cands)
        if d<=14: lab[dsu.find(key(p))].add(t)
    for c in gnd_components(segs,dsu,key): lab[c].add('GND')
    return page,segs,dsu,key,comp,lab

def probe(pgno, probes, maxd=15):
    page,segs,dsu,key,comp,lab=page_netmap(pgno)
    pts=[]
    for a,b in segs: pts.append(a); pts.append(b)
    for name,q in probes.items():
        d,p=min((math.dist(p,q),p) for p in pts)
        if d>maxd: print(f'{name:8s} -> NO WIRE'); continue
        c=dsu.find(key(p))
        print(f'{name:8s} d={d:.1f} n={len(comp[c]):3d}  NET = {sorted(lab.get(c,[])) or "(unlabeled)"}')