#!/usr/bin/env python3

import re
import time
import math
from dataclasses import replace
from itertools import chain
import subprocess

import gerbonara.cad.kicad.pcb as pcb
import gerbonara.cad.primitives as cad_pr
import gerbonara.cad.kicad.graphical_primitives as kc_gr


r = 85
edge_r = 95
cx, cy = 150, 100

prev_next_cycle = lambda l: zip(l[-1:] + l[:-1], l, l[1:] + l[:1])

si_prefixes = {'k': 1e3, 'M': 1e6, 'G': 1e9, 'm': 1e-3, 'u': 1e-6, 'µ': 1e-6, 'n': 1e-9, 'p': 1e-12, '.': 1, 'r': 1}
def parse_si(s, unit=''):
    if not (match := re.fullmatch(f'([0-9]*)(([.pPnNuUµmrRkKMgG])([0-9]*))?([pPnNuUµmrRkKMgG]?)({unit})?', s)):
        raise ValueError(f'{s} is not a valid number with an SI prefix')
    
    left, _1, dot, right, prefix, _unit = match.groups()
    prefix = prefix or dot or '.'
    multiplier = si_prefixes.get(prefix, si_prefixes.get(prefix.upper(), si_prefixes.get(prefix.lower())))

    return float(f'{left}.{right or ""}0') * multiplier
    

b = pcb.Board.open('self-balancing-test-a.kicad_pcb')
b.unfill_zones()
b.remove_many(chain(*(b.find_traces(net) for net in ('/MESH1A', '/MESH1B', '/MESH2A', '/MESH2B'))))

matches = list(b.find_footprints(name='Package_TO_SOT_SMD:SOT-23', sheetfile='resistor_bank.kicad_sch'))

components = []
for transistor in sorted(matches, key=lambda fp: fp.parsed_reference[1]):
    r1, r2 = transistor.pad(3).find_connected_footprints(name='Resistor_SMD.*')
    if parse_si(r1.value) > parse_si(r2.value):
        r1, r2 = r2, r1

    cap, = transistor.pad(3).find_connected_footprints(name='Capacitor_SMD.*')

    components.append((transistor, r1, r2, cap))

# Remove old generated traces and vias
b.remove_many(chain(
    [via for via in b.vias if not via.locked and not via.at.within_distance(cx, cy, r-20)],
    [seg for seg in b.track_segments if not seg.locked and not (seg.start.within_distance(cx, cy, r-20) or
                                                                seg.end.within_distance(cx, cy, r-20))],
    [arc for arc in b.track_arcs if not (arc.start.within_distance(cx, cy, r-20) or
                                         arc.end.within_distance(cx, cy, r-20))],
    ))

# Update git revision number in version silkscreen label
try:
    proc = subprocess.run(['git', 'describe', '--always'], check=True, capture_output=True, text=True)
    revision = proc.stdout.strip()
except subprocess.CalledProcessError as e:
    print('Warning: Cannot determine git revision', file=sys.stderr)
    print(e.stdout, file=sys.stderr)
    print(e.stderr, file=sys.stderr)
    revision = '<unknown rev>'

for t in b.texts:
    if t.text.startswith('Version'):
        version, _, _commit = t.text.partition(',')
        t.text = f'{version}, {revision}'
        t.render_cache = None

# Remove old, temporary board outline
b.circles = [c for c in b.circles if c.layer != 'Edge.Cuts']
    
# Generate new features
num_lgs = 8
lg_pitch = 1
segment_angle = 2*math.pi / len(matches)
print(f'Found {len(matches)} segments.')

for i, (fp, res1, res2, cap) in enumerate(components):
    alpha = i*segment_angle

    fp.at.x, fp.at.y = cx+r, cy
    fp.set_rotation(0)
    for text in fp.texts:
        text.at.rotation = 0
    for prop in fp.properties:
        prop.at.rotation = 0

    cap.at.x, cap.at.y  = cx+r+5, cy
    cap.side = 'back'
    cap.face('top', pad=1)

    res1.at.x, res1.at.y = cx+r+5, cy+2
    res1.face('right', net='*/V+')

    res2.at.x, res2.at.y = cx+r+5, cy-2
    res2.face('right', net='*/V+')

    gnd_via = pcb.Via(at=pcb.XYCoord(fp.pad(2)).with_offset(y=1.5),
                  size=1.2,
                  drill=0.6,
                  net=fp.pad(2).net.number
                  )
    b.add(gnd_via)

    lg = int(re.search(r'[0-9]*$', fp.pad(1).net.name).group(0))
    lg_via = pcb.Via(at=pcb.XYCoord(fp.pad(1)).with_offset(x=-2-lg_pitch*lg),
                  size=0.8,
                  drill=0.4,
                  net=fp.pad(1).net.number
                  )
    b.add(lg_via)

    x1, y1, _r, _f = res1.pad(2).abs_pos
    x2, y2, _r, _f = res2.pad(2).abs_pos
    x3, y3, _r, _f = res1.pad(1).abs_pos
    x4, y4, _r, _f = res2.pad(1).abs_pos
    mp = (x1+x2)/2, (y1+y2)/2
    for tr, net_name in [
        (cad_pr.Trace(1.5, res1.pad(2), res2.pad(2), orientation=['ccw']), fp.pad(3).net.name),
        (cad_pr.Trace(0.5, fp.pad(3), mp, orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, mp, (mp[0]+1.5, mp[1]), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x1, y1), (x1, y1+2), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x1-1, y1+2), (x1+1.6, y1+2), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x2, y2), (x2, y2-2), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x2-1, y2-2), (x2+1.6, y2-2), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x3, y3), (x3+1.5, y3), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x4, y4), (x4+1.5, y4), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x3+1.5, y3-0.8), (x3+1.5, y3+0.8), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x4+1.5, y4-0.8), (x4+1.5, y4+0.8), orientation=['cw']), fp.pad(3).net.name),
        (cad_pr.Trace(0.5, res1.pad(1), res2.pad(1)), res1.pad(1).net.name),
        (cad_pr.Trace(0.5, fp.pad(2), gnd_via), fp.pad(2).net.name),
        (cad_pr.Trace(0.5, fp.pad(1), lg_via), fp.pad(1).net.name),
        (cad_pr.Trace(1.5, (x1-1, y1+2), (x1+1.6, y1+2), orientation=['cw'], side='bottom'), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x2-1, y2-2), (x2+1.6, y2-2), orientation=['cw'], side='bottom'), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, mp, (mp[0]+1.5, mp[1]), orientation=['cw'], side='bottom'), fp.pad(3).net.name),
        (cad_pr.Trace(1.5, (x1+1.6, y1+2), (x2+1.6, y2-2), orientation=['cw'], side='bottom'), fp.pad(3).net.name),
        ]:
        for obj in b.map_gn_cad(tr, net_name=net_name):
            obj.rotate(alpha, cx, cy)
            b.add(obj)

    for j in range(-1, 2):
        for k in range(3):
            therm_via = pcb.Via(at=pcb.XYCoord(mp).with_offset(x=(k-1)*1.3 + 0.3, y=j*4),
                          size=1.2,
                          drill=0.6,
                          net=fp.pad(3).net.number
                          )
            b.add(therm_via)
            therm_via.rotate(alpha, cx, cy)
    
    for x, y in (x3, y3), (x4, y4):
        for k in -1, 1:
            therm_via = pcb.Via(at=(x+1.5, y+k*0.8),
                          size=1.2,
                          drill=0.6,
                          net=fp.pad(3).net.number
                          )
            b.add(therm_via)
            therm_via.rotate(alpha, cx, cy)

    for j in range(num_lgs):
        if 0 <= i-j < len(matches)-8:
            p1 = pcb.XYCoord(fp.pad(1)).with_offset(x=-2-lg_pitch*j)
            arc = pcb.TrackArc(
                    start=p1,
                    end=p1.with_rotation(segment_angle, cx, cy),
                    center=(cx, cy),
                    layer='B.Cu',
                    width=0.5,
                    net=b.net_id(f'LG{i}'),
                    )
            arc.rotate(alpha, cx, cy)
            b.add(arc)

    gnd_via.rotate(alpha, cx, cy)
    lg_via.rotate(alpha, cx, cy)
    fp.rotate(alpha, cx, cy)
    res1.rotate(alpha, cx, cy)
    res2.rotate(alpha, cx, cy)
    cap.rotate(alpha, cx, cy)

for i, ((prev_q, prev_r1, prev_r2, prev_cap), (fp, res1, res2, cap), (next_q, next_r1, next_r2, next_cap)) in enumerate(prev_next_cycle(components)):
    if res1.pad(1).net == next_r2.pad(1).net:
        arc = pcb.TrackArc(
                end=res1.pad(1),
                center=(cx, cy),
                start=next_r2.pad(1),
                width=0.5,
                net=res2.pad(1).net.number,
                )
        b.add(arc)

    # Add board outline segment
    segment_width = 10.0 # mm
    segment_length = 10.0 # mm
    fillet = 1.0 # mm

    fillet_center_y = segment_width/2 - fillet
    fillet_center_x = math.sqrt((edge_r - fillet)**2 - fillet_center_y**2)

    tangent_y = fillet_center_y * edge_r / fillet_center_x
    tangent_x = math.sqrt(edge_r**2 - tangent_y**2)
    
    inner_fillet_cos = fillet * math.cos(segment_angle/2)
    inner_fillet_sin = fillet * math.sin(segment_angle/2)

    c = pcb.XYCoord(cx, cy)

    arc = kc_gr.Arc(
            start=c.with_offset(tangent_x, -tangent_y),
            end=c.with_offset(tangent_x, tangent_y),
            center=(cx, cy),
            layer='Edge.Cuts',
            width=0.2)
    arc.rotate(i*segment_angle, cx, cy)
    b.add(arc)

    arc = kc_gr.Arc(
            start=c.with_offset(tangent_x, tangent_y),
            end=c.with_offset(fillet_center_x, segment_width/2),
            center=c.with_offset(fillet_center_x, fillet_center_y),
            layer='Edge.Cuts',
            width=0.2)
    arc.rotate(i*segment_angle, cx, cy)
    b.add(arc)

    arc = kc_gr.Arc(
            start=c.with_offset(fillet_center_x, -segment_width/2),
            end=c.with_offset(tangent_x, -tangent_y),
            center=c.with_offset(fillet_center_x, -fillet_center_y),
            layer='Edge.Cuts',
            width=0.2)
    arc.rotate(i*segment_angle, cx, cy)
    b.add(arc)

    l2 = kc_gr.Line(
            start=c.with_offset(fillet_center_x, segment_width/2),
            end=c.with_offset(edge_r - segment_length, segment_width/2),
            layer='Edge.Cuts',
            width=0.2)
    l2.rotate(i*segment_angle, cx, cy)
    b.add(l2)

    l2 = kc_gr.Line(
            start=c.with_offset(fillet_center_x, -segment_width/2),
            end=c.with_offset(edge_r - segment_length, -segment_width/2),
            layer='Edge.Cuts',
            width=0.2)
    l2.rotate(i*segment_angle, cx, cy)
    b.add(l2)

    p1 = c.with_offset(edge_r - segment_length, segment_width/2 + fillet)
    p2 = p1.with_offset(-inner_fillet_cos, -inner_fillet_sin)
    arc = kc_gr.Arc(
            start=pcb.XYCoord(p2),
            end=c.with_offset(edge_r - segment_length, segment_width/2),
            center=pcb.XYCoord(p1),
            layer='Edge.Cuts',
            width=0.2)
    arc.rotate(i*segment_angle, cx, cy)
    b.add(arc)

    p1 = c.with_offset(edge_r - segment_length, -segment_width/2 - fillet)
    p3 = p1.with_offset(-inner_fillet_cos, -inner_fillet_sin)
    arc = kc_gr.Arc(
            end=pcb.XYCoord(p3),
            start=c.with_offset(edge_r - segment_length, -segment_width/2),
            center=pcb.XYCoord(p1),
            layer='Edge.Cuts',
            width=0.2)
    arc.rotate(i*segment_angle, cx, cy)
    b.add(arc)

    l2 = kc_gr.Line(
            start=pcb.XYCoord(p3),
            end=p2.with_rotation(segment_angle, cx, cy),
            layer='Edge.Cuts',
            width=0.2)
    l2.rotate(i*segment_angle, cx, cy)
    b.add(l2)

b.write('self-balancing-test-a.testout.kicad_pcb')