from pathlib import Path
import wx,json,math,shutil
_app=wx.App(False)
import pcbnew as k
R=Path(__file__).resolve().parent;P=R/'board/led_driver.kicad_pcb';b=k.LoadBoard(str(P));shutil.copy2(P,R/'before_ground_stitch.kicad_pcb');polys={}
for z in b.Zones():
 for l in z.GetLayerSet().Seq():polys[l]=z.GetFilledPolysList(l)
F,B=k.F_Cu,k.B_Cu
parents={(l,i):(l,i) for l,p in polys.items() for i in range(p.OutlineCount())}
def root(a):
 while parents[a]!=a:a=parents[a]
 return a
def union(nodes):
 if not nodes:return
 a=root(nodes[0])
 for n in nodes[1:]:parents[root(n)]=a
def nodes(xy):
 p=k.VECTOR2I(k.FromMM(xy[0]),k.FromMM(xy[1]));return [(l,i) for l,poly in polys.items() for i in range(poly.OutlineCount()) if poly.Contains(p,i)]
holes=[]
for t in b.GetTracks():
 if isinstance(t,k.PCB_VIA):
  xy=(k.ToMM(t.GetPosition().x),k.ToMM(t.GetPosition().y));holes.append((*xy,k.ToMM(t.GetDrill())/2))
  if t.GetNetname()=='/GND':union(nodes(xy))
 else:
  if t.GetNetname()=='/GND':union([n for n in nodes((k.ToMM(t.GetStart().x),k.ToMM(t.GetStart().y)))+nodes((k.ToMM(t.GetEnd().x),k.ToMM(t.GetEnd().y))) if n[0]==t.GetLayer()])
for f in b.GetFootprints():
 for pad in f.Pads():
  xy=(k.ToMM(pad.GetPosition().x),k.ToMM(pad.GetPosition().y));dr=max(pad.GetDrillSize().x,pad.GetDrillSize().y)
  if dr:holes.append((*xy,k.ToMM(dr)/2))
  if dr and pad.GetNetname()=='/GND':union(nodes(xy))
before=len({root(n) for n in parents});adds=[]
for ix in range(8,189):
 for iy in range(8,177):
  xy=(ix*.5,iy*.5)
  if any(math.dist(xy,(x,y))<rad+.15+.26 for x,y,rad in holes):continue
  ns=nodes(xy)
  if len(ns)<2 or len({root(n) for n in ns})<2:continue
  if len({l for l,i in ns})<2:continue
  if not all(all(polys[l].Contains(k.VECTOR2I(k.FromMM(xy[0]+.325*math.cos(a*math.pi/4)),k.FromMM(xy[1]+.325*math.sin(a*math.pi/4))),i) for a in range(8)) for l,i in ns):continue
  t=k.PCB_VIA(b);t.SetPosition(k.VECTOR2I(k.FromMM(xy[0]),k.FromMM(xy[1])));t.SetWidth(k.FromMM(.6));t.SetDrill(k.FromMM(.3));t.SetViaType(k.VIATYPE_THROUGH);t.SetLayerPair(F,B);t.SetNetCode(b.FindNet('/GND').GetNetCode());b.Add(t);adds.append({'xy':xy,'joins':ns});union(ns);holes.append((*xy,.15))
  if len(adds)>=10:break
 if len(adds)>=10:break
groups={}
for n in parents:groups.setdefault(str(root(n)),[]).append(n)
print('GROUPS',json.dumps(groups))
for f in b.GetFootprints():
 for pad in f.Pads():
  if pad.GetNetname()=='/GND':
   xy=(k.ToMM(pad.GetPosition().x),k.ToMM(pad.GetPosition().y));print('PAD',f.GetReference(),pad.GetNumber(),xy,[str(root(n)) for n in nodes(xy)])
