distribute.py

#!/usr/bin/env python2

from os import sys
import argparse
import copy
import pcbnew
import string
import sys
import math
import numpy as np

parser = argparse.ArgumentParser(description='distribute leds along a line')

parser.add_argument('board', metavar='b', type=str,
        help='Input board')

args = parser.parse_args()

board = pcbnew.LoadBoard(args.board)

layertable = {}

numlayers = pcbnew.PCB_LAYER_ID_COUNT
for i in range(numlayers):
    layertable[board.GetLayerName(i)] = i

BFab = layertable['B.Fab']

tp_labels = []
mounting_holes = []

print("\nCreating output board")
output = pcbnew.BOARD()

def find_nearest(array,value):
    return np.argmin(np.abs(array - value))

board_width = board.ComputeBoundingBox().GetWidth()
def packpos(pos):
    return pos.y * board_width + pos.x

def unpackpos(pos):
    return pcbnew.wxPoint(pos % board_width, floor(pos / board_width))

def sortleds(s):
    return int(s.GetReference().strip(string.ascii_letters))

segments = []
leds = []

overall_length = 0
led_spacing = 0

for module in board.GetModules():
    if(module.GetValue() == "APA102"):
        leds.append(module)

for drawing in board.GetDrawings():
    if isinstance(drawing, pcbnew.DRAWSEGMENT) and drawing.GetLayer() == BFab:
        start = drawing.GetStart()
        end = drawing.GetEnd()

        overall_length += math.sqrt(math.pow(start.x - end.x, 2) + math.pow(start.y - end.y, 2))

        segments.append(drawing)

leds.sort(key=sortleds)

led_spacing = overall_length / len(leds) - 1
start_pos = leds[0].GetPosition()

start_seg = None
start_first = False
diff = -1

def seglen(seg):
    return math.sqrt(math.pow(seg.x, 2) + math.pow(seg.y, 2))

print("starting LED (%s) [ x: %f, y: %f ]" % (leds[0].GetReference(), start_pos.x, start_pos.y))
print("finding closest segment")

for seg in segments:
    start_diff = seglen(seg.GetStart() - start_pos)
    end_diff  = seglen(seg.GetEnd() - start_pos)

    if(diff == -1 or (start_diff < diff or end_diff < diff)):
        start_seg = seg
        if(start_diff < end_diff):
            start_first = True
            diff = start_diff
        else:
            start_first = False
            diff = end_diff

def swap_points(seg):
    start = seg.GetStart();
    startx = start.x
    starty = start.y

    end = seg.GetEnd();
    endx = end.x
    endy = end.y

    seg.SetStartX(endx)
    seg.SetStartY(endy)

    seg.SetEndX(startx)
    seg.SetEndY(starty)

print("closest segment point: %d" % (diff))
if start_first is False:
    swap_points(start_seg)

print("ordering segments from nearest")

def order_segments(start_seg):
    ordered_segs = [ start_seg ]

    next_seg = start_seg

    if(seglen(start_seg.GetStart() - start_pos) < seglen(start_seg.GetEnd() - start_pos)):
        next_pos = start_seg.GetEnd();
    else:
        next_pos = start_seg.GetStart();

    dist = -1

    while(len(segments) > 0):
        for seg in segments:
            start = seg.GetStart()
            end = seg.GetEnd()

            #print(next_pos==start,next_pos==end, next_pos, start, end)
            if(next_pos == start or next_pos == end):
                segments.remove(seg)
                ordered_segs.append(seg)
                next_seg = seg

                end = seg.GetEnd()

                if(next_pos == end):
                    swap_points(next_seg)

                next_pos = end
                break

    return ordered_segs

segments.remove(start_seg)
segments = order_segments(start_seg)

print("overall segment length: %d" % (overall_length))
print("led spacing: %d" % (led_spacing))

seg_idx = 0
offset = 0
needed = led_spacing

for i, led in enumerate(leds[1:], start=1):

    seg = segments[seg_idx]
    start = seg.GetStart()
    end = seg.GetEnd()
    v = end - start
    length = seglen(v)

    if(seg_idx + 2 > len(segments)): break

    print("\nProcessing LED: %d" % ( i ))
    print("Needed: %d length (seq: %d): %d offset %d" % ( needed, seg_idx, length, offset ))

    while(needed > (length - offset)):
        seg_idx += 1
        needed -= (length - offset)
        offset = 0

        seg = segments[seg_idx]

        start = seg.GetStart()
        end = seg.GetEnd()

        v = end - start
        length = seglen(v)

        print("Needed: %d length (seq: %d): %d offset %d" % ( needed, seg_idx, length, offset ))

    vx = end.x - start.x
    vy = end.y - start.y

    mag = math.sqrt(vx*vx + vy*vy)

    vx /= mag
    vy /= mag

    x = start.x + vx * (offset + needed)
    y = start.y + vy * (offset + needed)

    offset += led_spacing

    led.SetPosition(pcbnew.wxPoint(x, y))

board.Save("distributed-" + args.board)