"""Original editable teaching assembly; arbitrary millimetres, not a device design.
Run with Python, cadquery==2.7.0 and matplotlib. All outputs beside this file.
No stiffness, damping, material, load rating or manufacturing claim is implied.
"""
from pathlib import Path
import csv, json, hashlib
import cadquery as cq
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.patches import Patch

OUT = Path(__file__).resolve().parent
P = dict(base_width=130., depth=66., base_thickness=8., gap=92.,
         mass_width=112., mass_thickness=24., spring_offset=36.,
         coil_radius=9., wire_radius=1.6, coil_turns=7., damper_radius=7.,
         damper_height=49., rod_radius=2.5)
COLORS = dict(base='#5676ae', mass='#75d9ed', spring='#e6adff', damper='#ffcf78')

def build():
    p=P; bottom=p['base_thickness']; top=bottom+p['gap']
    parts=[]
    def add(name, group, shape, role):
        parts.append((name,group,shape,role))
    add('Moving base','base',cq.Workplane('XY').box(p['base_width'],p['depth'],bottom).translate((0,0,bottom/2)).val(),'Receives the imposed motion; connection details omitted.')
    add('Inert supported mass','mass',cq.Workplane('XY').box(p['mass_width'],p['depth']*.82,p['mass_thickness']).translate((0,0,top+p['mass_thickness']/2)).val(),'Rigid inert payload symbol; geometry does not define the simulated mass.')
    for side in (-1,1):
        helix=cq.Wire.makeHelix(p['gap']/p['coil_turns'],p['gap'],p['coil_radius'])
        spring=cq.Workplane('XZ').center(p['coil_radius'],0).circle(p['wire_radius']).sweep(cq.Workplane(obj=helix),isFrenet=True).val().translate((side*p['spring_offset'],0,bottom))
        add(('Left' if side<0 else 'Right')+' spring','spring',spring,'Elastic storage symbol; illustrative coil geometry is not sized for any stiffness or stress.')
    add('Damper body','damper',cq.Solid.makeCylinder(p['damper_radius'],p['damper_height'],cq.Vector(0,0,bottom)),'Dissipation symbol; no piston, seals or fluid circuit are modeled.')
    add('Damper rod','damper',cq.Solid.makeCylinder(p['rod_radius'],p['gap']-p['damper_height'],cq.Vector(0,0,bottom+p['damper_height'])),'Telescoping rod symbol; no stroke or force rating.')
    return parts

def render(parts,mobile=False):
    fig=plt.figure(figsize=(6,8) if mobile else (12,8),facecolor='#08162e')
    ax=fig.add_subplot(111,projection='3d',facecolor='#08162e')
    ax.set_proj_type('ortho')
    for name,group,shape,role in parts:
        vs,ts=shape.tessellate(.4,.18)
        xyz=[(v.x,v.y,v.z) for v in vs]
        ax.add_collection3d(Poly3DCollection([[xyz[i] for i in t] for t in ts],facecolors=COLORS[group],edgecolors=None,shade=True,lightsource=matplotlib.colors.LightSource(azdeg=315,altdeg=45)))
    ax.set(xlim=(-75,75),ylim=(-40,40),zlim=(0,132))
    ax.set_box_aspect((150,80,132)); ax.view_init(elev=18,azim=-67)
    ax.set_axis_off(); ax.set_position([.02,.20,.96,.63])
    fig.text(.06,.94,'SHOCK ISOLATION',color='white',fontsize=21 if mobile else 27,weight='bold')
    fig.text(.06,.89,'A mass, springs and a damper',color='#a9c9f6',fontsize=13 if mobile else 17)
    handles=[Patch(color=COLORS[g],label=t) for g,t in [('mass','Inert mass'),('spring','Springs: store energy'),('damper','Damper: dissipates energy'),('base','Moving base')]]
    fig.legend(handles=handles,loc='lower center',bbox_to_anchor=(.5,.10),ncol=1 if mobile else 2,frameon=False,labelcolor='white',fontsize=11)
    fig.text(.06,.045,'Conceptual CAD • arbitrary geometry\nNot a tested device or manufacturing drawing.',color='#bdcce5',fontsize=10)
    stem='shock-isolation-cad-mobile' if mobile else 'shock-isolation-cad-wide'
    fig.savefig(OUT/(stem+'.png'),dpi=180,facecolor=fig.get_facecolor())
    fig.savefig(OUT/(stem+'.pdf'),facecolor=fig.get_facecolor())
    if not mobile:
        for item in list(fig.texts)+list(fig.legends):
            item.remove()
        fig.set_size_inches(7,7)
        ax.set_position([0,0,1,1])
        fig.savefig(OUT/'shock-isolation-cad-hero.png',dpi=180,facecolor=fig.get_facecolor())
    plt.close(fig)

def main():
    parts=build(); assembly=cq.Assembly(name='Conceptual shock isolator')
    for name,group,shape,role in parts:
        assert shape.isValid() and shape.Volume()>0
        assembly.add(shape,name=name,color=cq.Color(COLORS[group]))
    target=OUT/'shock-isolation-concept.step'
    assembly.save(str(target))
    restored=cq.importers.importStep(str(target)).solids().vals()
    assert len(restored)==len(parts)
    assert all(s.isValid() and s.Volume()>0 for s in restored)
    original_volume=sum(s.Volume() for _,_,s,_ in parts)
    restored_volume=sum(s.Volume() for s in restored)
    assert abs(original_volume-restored_volume)/original_volume<1e-8
    with (OUT/'components.csv').open('w',newline='',encoding='utf-8') as f:
        writer=csv.writer(f,lineterminator='\n');writer.writerow(['component','group','role_and_limit','volume_mm3_for_geometry_check_only'])
        writer.writerows((n,g,r,format(s.Volume(),'.12g')) for n,g,s,r in parts)
    (OUT/'validation.json').write_text(json.dumps(dict(cadquery_version=cq.__version__,geometry_parameters=P,parameter_units="Millimetres except dimensionless coil_turns",original_solids=len(parts),reimported_solids=len(restored),all_valid=True,positive_volume=True,original_volume_mm3=original_volume,reimported_volume_mm3=restored_volume,step_sha256=hashlib.sha256(target.read_bytes()).hexdigest(),scope='Geometry and round-trip validation only. No mechanical performance validation.'),indent=2)+'\n',encoding='utf-8')
    render(parts);render(parts,True)
    print('Validated',len(restored),'STEP solids and generated two views.')

if __name__=='__main__': main()