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

def highlight(pgno, probe, out, pad=22, scale=6):
    doc=fitz.open('Xiaomi_Mi_8_Schematic.pdf')
    page=doc[pgno-1]
    segs,dsu,key,comp=nets(page)
    pts=[]
    for a,b in segs: pts.append(a); pts.append(b)
    d,p=min((math.dist(p,probe),p) for p in pts)
    c=dsu.find(key(p))
    mine=[(a,b) for a,b in segs if dsu.find(key(a))==c]
    r=fitz.Rect(1e9,1e9,-1e9,-1e9)
    for a,b in mine:
        for q in (a,b):
            r.x0=min(r.x0,q[0]); r.y0=min(r.y0,q[1])
            r.x1=max(r.x1,q[0]); r.y1=max(r.y1,q[1])
    for a,b in mine:
        page.draw_line(fitz.Point(*a), fitz.Point(*b), color=(0,0.75,0), width=2.2)
    r=fitz.Rect(r.x0-pad, r.y0-pad, r.x1+pad, r.y1+pad)
    clip=r*page.rotation_matrix
    pix=page.get_pixmap(matrix=fitz.Matrix(scale,scale), clip=clip)
    pix.save(out)
    print(out, 'dist=%.1f'%d, 'segs=%d'%len(mine), 'bbox=%s'%r, pix.width, pix.height)