"""Size-constrained idealized bottle; no use of the corrupted radial profile."""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
from pathlib import Path
import json,xml.etree.ElementTree as E
import numpy as np
import trimesh,mujoco
from PIL import Image

OUT=Path(__file__).resolve().parents[1]/'third_party/bottle_model_parametric'
OUT.mkdir(parents=True,exist_ok=True);(OUT/'assets').mkdir(exist_ok=True)
H=.157856;R=.05411/2;neck=.014
# Only H and 2R are inherited dimensions. Shoulder/neck proportions are design choices.
profile=np.array([[0,R],[.103,R],[.139,neck],[H,neck]])
def lathe(points):
    n=128;a=np.arange(n)*2*np.pi/n
    vertices=np.concatenate([np.c_[r*np.cos(a),r*np.sin(a),np.full(n,z)] for z,r in points])
    faces=[]
    for k in range(len(points)-1):
        for i in range(n):
            j=(i+1)%n;faces.extend([[k*n+i,k*n+j,(k+1)*n+j],[k*n+i,(k+1)*n+j,(k+1)*n+i]])
    bottom=len(vertices);top=bottom+1
    vertices=np.r_[vertices,[[0,0,points[0,0]],[0,0,points[-1,0]]]]
    for i in range(n):
        j=(i+1)%n;faces.extend([[bottom,j,i],[top,(len(points)-1)*n+i,(len(points)-1)*n+j]])
    return trimesh.Trimesh(vertices=vertices,faces=faces,process=True)
mesh=lathe(profile);mesh.export(OUT/'assets/bottle.obj');mesh.export(OUT/'bottle.stl');mesh.export(OUT/'bottle.glb')
shoulder=lathe(profile[1:3]);shoulder.export(OUT/'assets/shoulder.obj')
root=E.Element('mujoco',model='dimension_based_bottle');E.SubElement(root,'option',gravity='0 0 -9.81')
asset=E.SubElement(root,'asset')
E.SubElement(asset,'mesh',name='bottle_visual',file='assets/bottle.obj')
E.SubElement(asset,'mesh',name='bottle_shoulder',file='assets/shoulder.obj')
world=E.SubElement(root,'worldbody');body=E.SubElement(world,'body',name='bottle')
E.SubElement(body,'freejoint',name='bottle_freejoint')
mesh.density=.2308/mesh.volume
E.SubElement(body,'inertial',mass='.2308',pos=' '.join(map(str,mesh.center_mass)),diaginertia=' '.join(map(str,np.diag(mesh.moment_inertia))))
E.SubElement(body,'geom',type='mesh',mesh='bottle_visual',group='1',contype='0',conaffinity='0',rgba='.68 .74 .8 1')
common=dict(group='3',friction='.6 .005 .001',condim='6')
E.SubElement(body,'geom',type='cylinder',size=f'{R} {.103/2}',pos=f'0 0 {.103/2}',**common)
E.SubElement(body,'geom',type='mesh',mesh='bottle_shoulder',**common)
E.SubElement(body,'geom',type='cylinder',size=f'{neck} {(H-.139)/2}',pos=f'0 0 {(H+.139)/2}',**common)
E.indent(root);E.ElementTree(root).write(OUT/'bottle.xml')
np.savez(OUT/'profile.npz',h=profile[:,0],r=profile[:,1],height=H,diameter=2*R)
model=mujoco.MjModel.from_xml_path(str(OUT/'bottle.xml'));data=mujoco.MjData(model);mujoco.mj_forward(model,data)
renderer=mujoco.Renderer(model,height=480,width=640);cam=mujoco.MjvCamera();cam.lookat[:]=[0,0,H/2];cam.distance=.36;cam.azimuth=60;cam.elevation=-8
opt=mujoco.MjvOption();opt.geomgroup[3]=0;model.vis.headlight.ambient[:]=.7
renderer.update_scene(data,camera=cam,scene_option=opt);Image.fromarray(renderer.render()).save(OUT/'preview.png');renderer.close()
report=dict(height_mm=H*1000,max_diameter_mm=2*R*1000,neck_diameter_mm=neck*2000,watertight=bool(mesh.is_watertight),winding_consistent=bool(mesh.is_winding_consistent),volume_m3=mesh.volume,mujoco_load_passed=True,collision_parts=3,extent_m=mesh.extents.tolist(),shape_source='Idealized cylinder body, tapered shoulder, cylinder neck; only height and maximum diameter inherited',mass_source='Previous XML mass 0.2308 kg retained; uniform solid inertia recomputed, not measured')
(OUT/'validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2))