"""Rwanda pen plotter: original parametric CAD inspired by the Oulu printable axis.
All dimensions are mm. Hardware is reference geometry, NOT printable substitutes.
Build the selected machine: python build_cnc_option.py
This module exports shared printed motion parts; the CNC builder adds the plywood and adapters.
The default export directory is this file's parent directory's parent.
No upstream mesh or FreeCAD geometry is copied into this design.
"""
from pathlib import Path
import argparse, csv, json, math, itertools
from collections import Counter
import cadquery as cq

HERE=Path(__file__).resolve().parent
P=json.loads((HERE/'parameters.json').read_text())
PARTS={};INSTANCES=[]
COLORS={'frame':(0.12,0.21,0.28),'rail':(0.13,0.55,0.55),'carriage':(0.96,0.47,0.20),'wheel':(0.18,0.22,0.25),'bed':(0.80,0.84,0.82),'hardware':(0.46,0.49,0.52),'cord':(0.83,0.65,0.23),'pen':(0.08,0.12,0.17),'paper':(0.99,0.98,0.93)}

def box(a,b,c,x=0,y=0,z=0):return cq.Workplane('XY').box(a,b,c).translate((x,y,z))
def cyl(d,h,x=0,y=0,z=0):return cq.Workplane('XY').circle(d/2).extrude(h).translate((x,y,z))
def holes(obj,points,d=3.4,z=-20,h=200):
 for x,y in points:obj=obj.cut(cyl(d,h,x,y,z))
 return obj

def register(name,obj,category='frame',orientation=None):
 # Export one printable body per file, with a deliberate print orientation.
 assert obj.val().isValid(),name
 assert len(obj.solids().vals())==1,(name,'disconnected bodies',len(obj.solids().vals()))
 PARTS[name]={'shape':obj,'category':category,'orientation':orientation}
 return name

def add(name,xyz=(0,0,0),rz=0,moving=None,kind='printed',shape=None,color=None,label=None):
 obj=shape or PARTS[name]['shape']
 if rz:obj=obj.rotate((0,0,0),(0,0,1),rz)
 obj=obj.translate(xyz)
 cat=PARTS[name]['category'] if kind=='printed' else kind
 INSTANCES.append({'name':f'{name}_{len(INSTANCES):03}', 'part':name,'shape':obj,'kind':kind,'moving':moving,'color':color or COLORS.get(cat,COLORS['hardware']),'label':label or name.replace('_',' ')})

def rail(length):
 # V faces are 45 degrees. The central spine stiffens the segmented beam.
 profile=[(-10,0),(10,0),(10,3),(15,8),(10,13),(8,13),(8,20),(-8,20),(-8,13),(-10,13),(-15,8),(-10,3)]
 obj=cq.Workplane('YZ').polyline(profile).close().extrude(length).translate((-length/2,0,0))
 pts=[(u,0) for u in [-length/2+15,-length/2+35,-26,-16,16,26,length/2-35,length/2-15]]
 obj=holes(obj,pts)
 for x,y in pts:obj=obj.cut(cyl(6.4,4.1,x,y,16))
 return obj

def axis_parts(prefix,length):
 obj=rail(length)
 for side in [-1,1]:
  part=obj.intersect(box(length/2-0.1,60,60,side*(length/4+0.05),0,10))
  register(f'{prefix}_rail_{"left" if side<0 else "right"}',part,'rail')

axis_parts('x',P['x_rail_length']);axis_parts('y',P['y_rail_length'])
register('rail_joiner',holes(box(66,20,6,z=-3),[(-26,0),(-16,0),(16,0),(26,0)]))
# End brackets bolt below the rail. NEMA 17 mounting pitch: 31 mm.
motor=box(50,100,4,y=30,z=-2)
motor=holes(motor,[(-10,0),(10,0)])
motor=holes(motor,[(u,46+v) for u in [-15.5,15.5] for v in [-15.5,15.5]],3.4)
motor=motor.cut(cyl(22.5,10,0,46,-5))
register('motor_bracket',motor)
idler=holes(box(50,94,4,y=27,z=-2),[(-10,0),(10,0)])
idler=idler.union(cyl(10,2,0,46,0))
idler=holes(idler,[(0,46)],5.3)
# Adjustment slot moves the bracket on the rail; two rail attachment holes are slots.
for u in [-10,10]:idler=idler.cut(cq.Workplane('XY').center(u,0).slot2D(8,3.4,0).extrude(10,both=True))
register('idler_bracket',idler)
# Two bearings per wheel with an internal spacer between inner races.
wheel=cq.Workplane('XZ').polyline([(3.9,0),(13,0),(13,2.5),(9.5,6),(13,9.5),(13,12),(3.9,12)]).close().revolve(360,(0,0),(0,1))
wheel=wheel.cut(cyl(P['wheel_bearing_seat_diameter'],P['wheel_bearing_seat_depth'],z=-0.01)).cut(cyl(P['wheel_bearing_seat_diameter'],P['wheel_bearing_seat_depth']+0.01,z=12-P['wheel_bearing_seat_depth']))
register('v_wheel',wheel,'wheel')
register('bearing_inner_spacer',cyl(7.5,1.8).cut(cyl(5.3,3,z=-.5)),'wheel')
register('axle_spacer',cyl(8,8).cut(cyl(5.3,10,z=-1)),'wheel')
bush=cyl(7,6).union(cyl(12,1.8,z=6)).cut(cyl(5.3,10,x=.6,z=-1))
register('eccentric_bush',bush,'carriage')
# Universal carriage with two screw-and-washer cord anchors at different heights.
v=P['wheel_center_offset'];b=P['drive_offset']-(P['capstan_diameter']+P['cord_diameter'])/2
car=box(66,72,6,z=25).union(box(24,24,6,y=38,z=25))
car=holes(car,[(u,v) for u in [-24,24]],5.3)
car=holes(car,[(u,-v) for u in [-24,24]],7.2)
car=holes(car,[(u,w) for u in [-16,16] for w in [-12,12]])
car=holes(car,[(-14,-28),(14,-28)])
car=car.union(box(24,16,3,y=40,z=1.5)).union(box(28,3,22,y=48.5,z=14))
for u,top in [(-8,P['cord_entry_height']-.3),(8,P['cord_return_height']-.3)]:
 car=car.union(cyl(8,top-3,u,b,3))
car=holes(car,[(-8,b),(8,b)])
register('carriage',car,'carriage')
# Smooth capstan, retained with an M3 set screw and a captured nut below the winding band.
capstan=cyl(P['capstan_diameter'],15,z=2).union(cyl(25,1.5,z=.5)).union(cyl(25,1.5,z=17)).union(cyl(12,2,z=-1.5))
capstan=capstan.cut(cyl(5.2,24,z=-3))
radial=cyl(3.4,15).rotate((0,0,0),(0,1,0),90).translate((0,0,3.5))
nut=cq.Workplane('YZ').polygon(6,6.8).extrude(3).translate((3.5,0,3.5))
# Top-loading nut channel stays inside the drum's winding surface.
nut_access=box(3,6.8,16,x=5,y=0,z=11.5)
capstan=capstan.cut(radial).cut(nut).cut(nut_access)
register('motor_capstan',capstan,'carriage')
# Far-end roller with two 625 bearings; the upper cord run returns at Z13.
roller=cyl(20,12,z=5).union(cyl(25,1.5,z=3.5)).union(cyl(25,1.5,z=17))
roller=roller.cut(cyl(7.8,20,z=1)).cut(cyl(P['wheel_bearing_seat_diameter'],5.11,z=4.99)).cut(cyl(P['wheel_bearing_seat_diameter'],5.11,z=11.9))
register('cord_idler',roller,'wheel')
register('idler_inner_support',cyl(7.5,3,z=2).cut(cyl(5.3,5,z=1)),'wheel')
# Frame: three side pieces per side, all below 190 mm for a 200 mm bed.
# Offset seams avoid the accessible X-tower base bolts at Y23/Y83.
side_points=[(0,y) for seam in [-90,95] for y in [seam-24,seam-10,seam+10,seam+24]]
side_points += [(0,23),(0,83)]+[(x,y) for x in [-8,8] for y in [-185,0,185]]
for name,a,z in [('base_side_front',-200,-90),('base_side_middle',-90,95),('base_side_rear',95,200)]:
 side=holes(box(26,z-a,10,y=(a+z)/2,z=5),side_points,4.4)
 register(name,side)
cross=holes(box(150,26,10,x=75,z=15),[(x,0) for x in [10,24,142]],4.4)
register('base_cross_half',cross)
strap=holes(box(60,24,5,z=2.5),[(x,0) for x in [-24,-10,10,24]],4.4)
register('base_join_strap',strap)
# Tall X supports bolt to the side bars; end rail bracket rests on their cap.
tower=box(46,72,8,z=14).union(box(20,44,108,z=72))
tower=tower.union(box(46,48,8,z=122))
# Side gussets are modelled as triangular prisms, not decorative disconnected ribs.
gusset=cq.Workplane('YZ').polyline([(-30,17),(18,17),(18,116)]).close().extrude(5)
tower=tower.union(gusset.translate((-14,0,0))).union(gusset.translate((9,0,0)))
tower=holes(tower,[(0,-30),(0,30)],4.4)
# Cap hardware has matching holes on a short clip plate under the rail extremity.
tower=holes(tower,[(-10,0),(10,0)],3.4)
for u in [-10,10]:
 # Side-access nut recesses below the rail-support cap.
 tower=tower.cut(cq.Workplane('XY').polygon(6,6.8).extrude(3.5).translate((u,0,114.5)))
register('x_tower',tower)
# Y end feet stand above each crossbar splice and support the end bracket.
foot=box(36,70,8,y=15,z=29).union(box(24,24,9,y=-15,z=37.5)).union(box(36,60,4,y=-25,z=44))
foot=holes(foot,[(-10,25),(10,25)],4.4)
foot=holes(foot,[(0,-15),(0,-35)],3.4)
register('y_end_foot',foot)
centerfoot=holes(box(32,64,19,z=34.5),[(0,y) for y in [-26,-16,16,26]],3.4)
centerfoot=holes(centerfoot,[(-10,0),(10,0)],4.4)
for u in [-10,10]:centerfoot=centerfoot.cut(cyl(8,4.21,u,0,39.8))
register('y_center_foot',centerfoot)
# Printed bed: four quadrants with integrated mounting bosses and underside ribs.
for sx,sy in itertools.product([-1,1],repeat=2):
 tile=box(120,95,4,sx*60,sy*47.5,98)
 tile=tile.union(box(12,12,18,sx*12,sy*16,87))
 # The ribs stop short of each joint strap and the carriage mounting region.
 tile=tile.union(box(4,62,6,sx*110,sy*51,93)).union(box(88,4,6,sx*64,sy*85,93))
 tile=tile.union(box(80,4,6,sx*68,sy*26,93))
 mount=[(sx*12,sy*16),(sx*60,sy*9),(sx*9,sy*45)]
 tile=holes(tile,mount)
 for x,y in mount:tile=tile.cut(cyl(6.4,2.3,x,y,97.7))
 register(f'bed_{"left" if sx<0 else "right"}_{"front" if sy<0 else "rear"}',tile,'bed')
bedstrap=holes(box(30,36,4,z=94),[(0,-9),(0,9)])
register('bed_join_strap',bedstrap,'bed')
# Fixed pen-stage bracket and 45-degree printed dovetail guides.
stage=box(36,24,6,y=-35,z=31).union(box(40,4.5,56,y=-40.75,z=0))
for u in [-12,12]:
 rib=cq.Workplane('XY').polyline([(u-3,-43),(u+3,-43),(u+5,-47),(u-5,-47)]).close().extrude(38).translate((0,0,-24))
 stage=stage.union(rib)
stage=holes(stage,[(-14,-28),(14,-28)])
# Lateral platform for a side-mounted micro servo, held by two cable-tie slots.
stage=stage.union(box(14,30,4,x=-39,y=-36,z=18)).union(box(28,16,6,x=-28,y=-23,z=34))
stage=stage.union(box(6,10,11,x=-38,y=-30,z=25.5))
for y in [-46,-23]:stage=stage.cut(box(9,2.5,8,-39,y,18))
register('pen_stage',stage,'carriage')
clear=P['slider_clearance']
slider=box(40,7.7,18,y=-47.15,z=-11)
for u in [-12,12]:
 cutter=cq.Workplane('XY').polyline([(u-3-clear,-42.5),(u+3+clear,-42.5),(u+5+clear,-47-clear),(u-5-clear,-47-clear)]).close().extrude(30).translate((0,0,-26))
 slider=slider.cut(cutter)
slider=slider.union(cyl(18,12,0,-53,-20)).cut(cyl(P['pen_diameter'],20,0,-53,-24))
slider=slider.cut(box(1.4,12,15,0,-64,-14))
slider=slider.union(box(7,7,12,-4.2,-62,-14)).union(box(7,7,12,4.2,-62,-14))
# Clamp screw is transverse, so it closes the split instead of piercing the pen.
clamp_hole=cyl(3.4,24).rotate((0,0,0),(0,1,0),90).translate((-12,-62,-14))
slider=slider.cut(clamp_hole)
slider=slider.union(box(12,6,3,-22,-52,-18.5))
register('pen_slider',slider,'carriage')
# Rail/seat/clearance coupons are optional prints, not included in the machine BOM.
coupon=box(76,28,6,z=3)
for i,d in enumerate([15.9,16.0,16.1,16.2]):coupon=coupon.cut(cyl(d,8,-27+i*18,0,-1))
register('bearing_fit_coupon',coupon,'bed')


def construct(x=0,y=0,pen_up=0):
 INSTANCES.clear()
 # Frame bars and seam straps.
 for sign in [-1,1]:
  for part in ['base_side_front','base_side_middle','base_side_rear']:add(part,(sign*150,0,0))
  for seam in [-90,95]:add('base_join_strap',(sign*150,seam,-5),90)
 for yy in [-185,0,185]:
  for side in [-1,1]:add('base_cross_half',(0,yy,0),0 if side>0 else 180)
  add('base_join_strap',(0,yy,20),0)
 # Supports.
 for xx in [-150,150]:add('x_tower',(xx,53,0))
 for sign in [-1,1]:add('y_end_foot',(0,sign*160,0),0 if sign>0 else 180)
 add('y_center_foot')
 # Two linear axes, in their local coordinates, with shared parts.
 for prefix,length,origin,angle,position in [('x',P['x_rail_length'],(0,53,P['x_axis_height']),0,x),('y',P['y_rail_length'],(0,0,P['y_axis_height']),90,y)]:
  def place(name,offset=(0,0,0),moving=None,kind='printed',shape=None,color=None):
   u,w,z=offset;c=math.cos(math.radians(angle));s=math.sin(math.radians(angle))
   xyz=(origin[0]+u*c-w*s,origin[1]+u*s+w*c,origin[2]+z)
   add(name,xyz,angle,moving,kind,shape,color)
  for side in ['left','right']:place(f'{prefix}_rail_{side}')
  place('rail_joiner')
  for part,sign in [('motor_bracket',-1),('idler_bracket',1)]:place(part,(sign*(length/2-25),0,0))
  place('carriage',(position,0,0),prefix)
  for u,w in itertools.product([-24,24],[-v,v]):
   bush_u=u
   if w<0:u+=.6
   place('v_wheel',(position+u,w,2),prefix)
   place('axle_spacer',(position+u,w,14),prefix)
   place('bearing_inner_spacer',(position+u,w,7.1),prefix)
   if w<0:place('eccentric_bush',(position+bush_u,w,22),prefix)
   bearing=cyl(16,5).cut(cyl(5,7,z=-1))
   for zz in [2.05,8.95]:place('625_bearing',(position+u,w,zz),prefix,'hardware',bearing)
  # Motor and bearing reference envelopes, plus the actual cord route.
  motorbody=box(42,42,40,z=-24).union(cyl(22,2,z=-4)).union(cyl(5,24,z=-4))
  end=length/2-25;offset=P['drive_offset'];r=(P['capstan_diameter']+P['cord_diameter'])/2
  place('nema17_motor',(-end,offset,0),kind='hardware',shape=motorbody)
  place('motor_capstan',(-end,offset,0))
  place('cord_idler',(end,offset,0))
  place('idler_inner_support',(end,offset,0))
  bearing=cyl(16,5).cut(cyl(5,7,z=-1))
  for zz in [5.05,11.95]:place('625_idler_bearing',(end,offset,zz),kind='hardware',shape=bearing)
  place('bearing_inner_spacer',(end,offset,10.1))
  lo=P['cord_entry_height'];hi=P['cord_return_height'];diam=P['cord_diameter']
  def cord_between(a,b):
   start=cq.Vector(*a);delta=cq.Vector(*b)-start
   return cq.Workplane(obj=cq.Solid.makeCylinder(diam/2,delta.Length,start,delta.normalized()))
  # Anchor -> lower motor tangent -> 3.5 helical wraps -> upper return roller -> anchor.
  place('cord_incoming',kind='cord',shape=cord_between((-end,offset-r,lo),(position-8,offset-r,lo)))
  INSTANCES[-1].update(cord_axis=prefix,cord_end=end,cord_segment='incoming')
  points=[]
  for i in range(253):
   t=i/252;theta=-math.pi/2-2*math.pi*P['capstan_wraps']*t
   points.append(cq.Vector(-end+r*math.cos(theta),offset+r*math.sin(theta),lo+(hi-lo)*t))
  path=cq.Workplane().spline(points)
  plane=cq.Plane(origin=points[0],xDir=(0,1,0),normal=(points[1]-points[0]).normalized())
  winding=cq.Workplane(plane).circle(diam/2).sweep(path,isFrenet=True)
  place('cord_winding',kind='cord',shape=winding)
  place('cord_outgoing',kind='cord',shape=cord_between((-end,offset+r,hi),(end,offset+r,hi)))
  points=[cq.Vector(end+r*math.cos(math.pi/2-math.pi*i/60),offset+r*math.sin(math.pi/2-math.pi*i/60),hi) for i in range(61)]
  plane=cq.Plane(origin=points[0],xDir=(0,0,1),normal=(points[1]-points[0]).normalized())
  turn=cq.Workplane(plane).circle(diam/2).sweep(cq.Workplane().spline(points),isFrenet=True)
  place('cord_return_turn',kind='cord',shape=turn)
  place('cord_to_anchor',kind='cord',shape=cord_between((position+8,offset-r,hi),(end,offset-r,hi)))
  INSTANCES[-1].update(cord_axis=prefix,cord_end=end,cord_segment='return')
 # Y bed moves with the paper; software must account for this sign.
 for sx,sy in itertools.product([-1,1],repeat=2):add(f'bed_{"left" if sx<0 else "right"}_{"front" if sy<0 else "rear"}',(0,y,0),moving='y')
 for xx in [-60,60]:add('bed_join_strap',(xx,y,0),moving='y')
 for yy in [-45,45]:add('bed_join_strap',(0,y+yy,0),90,moving='y')
 add('pen_stage',(x,53,130),moving='x')
 add('pen_slider',(x,53,130+pen_up),moving='x')
 pen=cyl(9.8,88,z=P['paper_height']+4).union(cq.Workplane('XY').circle(.5).workplane(offset=4).circle(4.9).loft().translate((0,0,P['paper_height'])))
 add('pen',(x,0,pen_up),moving='x',kind='pen',shape=pen)
 servo=box(12.2,23,27,-39,17,2.5).union(cyl(5,5).rotate((0,0,0),(0,1,0),90).translate((-33,8,-12)))
 add('sg90_servo',(x,0,130),moving='x',kind='hardware',shape=servo,color=(0.18,.36,.74))
 # Supplied horn envelope, contacting the back edge of the pickup tab.
 horn=cq.Workplane('YZ').polyline([(7,119),(9,117),(6,108),(4,110)]).close().extrude(3).translate((-28,0,0))
 horn=horn.union(cyl(6,3).rotate((0,0,0),(0,1,0),90).translate((-28,8,118)))
 add('servo_horn',(x,0,0),moving='x',kind='hardware',shape=horn,color=(.86,.88,.90))
 add('paper',(0,y,P['paper_height']+.1),moving='y',kind='paper',shape=box(220,170,.18))
 return INSTANCES


def export_shared_parts(out):
 import trimesh
 for d in ['stl','step','docs']: (out/d).mkdir(parents=True,exist_ok=True)
 construct()
 counts=Counter(i['part'] for i in INSTANCES if i['kind']=='printed')
 report={'units':'mm','status':'CAD prototype; not physically fabricated or calibrated','parameters':P,'printer_bed_mm':[200,200],'allowed_xy_mm':[190,190],'parts':[]}
 for name,part in PARTS.items():
  if name in {'bed_right_front', 'base_side_middle', 'bed_right_rear', 'base_side_rear', 'bed_join_strap', 'base_join_strap', 'base_side_front', 'bed_left_rear', 'base_cross_half', 'bed_left_front'} or name.startswith('cnc_'):continue
  obj=part['shape']
  # Export in stable, flat print coordinates. Tall parts are tipped on their side.
  if name=='pen_stage':obj=obj.rotate((0,0,0),(1,0,0),90)
  if name in ['carriage','motor_capstan'] or name.startswith('bed_'):obj=obj.rotate((0,0,0),(1,0,0),180)
  bb=obj.val().BoundingBox();obj=obj.translate((-bb.xmin,-bb.ymin,-bb.zmin))
  bb=obj.val().BoundingBox()
  cq.exporters.export(obj,str(out/'stl'/f'{name}.stl'),tolerance=.08,angularTolerance=.12)
  cq.exporters.export(obj,str(out/'step'/f'{name}.step'))
  mesh=trimesh.load_mesh(out/'stl'/f'{name}.stl')
  info={'part':name,'quantity':counts[name],'optional':name.endswith('coupon'),'solid_count':len(obj.solids().vals()),'valid':obj.val().isValid(),'watertight':bool(mesh.is_watertight),'size_mm':[round(n,2) for n in [bb.xlen,bb.ylen,bb.zlen]],'cad_volume_mm3':round(obj.val().Volume(),2)}
  assert info['watertight'],info
  assert max(info['size_mm'][:2])<=190 and info['size_mm'][2]<=200,info
  report['parts'].append(info)
  print('exported',name,info['size_mm'],flush=True)
 report['check_scope']='Individual shared-component BRep/STL checks; see cnc/cad-validation.json for the assembled plywood machine checks.'
 report['calculated_starting_steps_per_mm']=P['motor_steps_per_turn']*P['microsteps']/(math.pi*(P['capstan_diameter']+P['cord_diameter']))
 report['pen_up_clearance_mm3']=round(PARTS['pen_stage']['shape'].intersect(PARTS['pen_slider']['shape'].translate((0,0,6))).val().Volume(),6)
 (out/'docs'/'cad-validation.json').write_text(json.dumps(report,indent=2)+'\n')
 print('Exported',len(report['parts']),'shared printed part types.',flush=True)
 return report

if __name__=='__main__':
 parser=argparse.ArgumentParser(description='Export shared components for the plywood plotter; use build_cnc_option.py for the complete machine.')
 parser.add_argument('--out',type=Path,default=HERE.parent)
 args=parser.parse_args();export_shared_parts(args.out)