#!/usr/bin/env python3
"""Simulate and illustrate a 50-bean grinder-hopper discharge in MuJoCo."""
from pathlib import Path
import os
import math
import numpy as np
import mujoco
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from matplotlib.collections import LineCollection
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

OUT = Path(os.environ.get("OUT", "hopper.png"))
XML_PATH = Path(__file__).resolve().with_name("model.xml")
SEED = 17
N_BEANS = 50
DT = 0.0005
END = 5.0
THROAT_R = 0.009                 # 18 mm clear diameter
GATE_Z = 0.0630
COUNT_Z = 0.0605
CONE_Z0, CONE_Z1 = 0.080, 0.180
CONE_R1 = 0.047

# Closed, thick conical/cylindrical shells (triangles, not a visual-only cone).
def shell_mesh(name, r0, r1, z0, z1, thick=0.0012, n=64):
    verts = []
    for r, z in [(r0,z0),(r1,z1),(r1+thick,z1),(r0+thick,z0)]:
        verts.extend([(r*math.cos(2*math.pi*i/n), r*math.sin(2*math.pi*i/n), z) for i in range(n)])
    faces=[]
    def quad(a,b,c,d): faces.extend([(a,b,c),(a,c,d)])
    for i in range(n):
        j=(i+1)%n
        quad(i,n+i,n+j,j)                         # inner wall
        quad(3*n+i,3*n+j,2*n+j,2*n+i)             # outer wall
        quad(n+i,2*n+i,2*n+j,n+j)                 # top annulus
        quad(i,j,3*n+j,3*n+i)                     # bottom annulus
    v=" ".join(f"{x:.7f} {y:.7f} {z:.7f}" for x,y,z in verts)
    f=" ".join(str(k) for tri in faces for k in tri)
    return f'<mesh name="{name}" vertex="{v}" face="{f}"/>'

def initial_beans():
    rng=np.random.default_rng(SEED)
    counts=[3,7,10,13,17]
    zs=[0.096,0.114,0.132,0.150,0.168]
    result=[]
    for layer,(num,z) in enumerate(zip(counts,zs)):
        rin=THROAT_R+(CONE_R1-THROAT_R)*(z-CONE_Z0)/(CONE_Z1-CONE_Z0)
        rad=max(0.0,rin-0.0062)
        for j in range(num):
            # One central bean, remaining beans equally spaced on a jittered ring.
            rr=0 if (j==0 and num%2==1) else rad*(0.82 if num>7 else 0.70)
            a=2*math.pi*j/num + layer*0.37 + rng.uniform(-0.035,0.035)
            p=[rr*math.cos(a),rr*math.sin(a),z+rng.uniform(-0.0006,0.0006)]
            q=rng.normal(size=4); q/=np.linalg.norm(q)
            result.append((p,q))
    return result

cone_mesh=shell_mesh("hopper_shell",THROAT_R,CONE_R1,CONE_Z0,CONE_Z1)
tube_mesh=shell_mesh("throat_shell",THROAT_R,THROAT_R,0.0655,0.081)

def mat_to_quat(m):
    # MuJoCo's wxyz convention.
    tr=np.trace(m)
    if tr>0:
        s=math.sqrt(tr+1)*2; w=s/4; x=(m[2,1]-m[1,2])/s; y=(m[0,2]-m[2,0])/s; z=(m[1,0]-m[0,1])/s
    else:
        i=int(np.argmax(np.diag(m))); j=(i+1)%3; k=(i+2)%3
        s=math.sqrt(1+m[i,i]-m[j,j]-m[k,k])*2
        q=[0.,0.,0.,0.]; q[0]=(m[k,j]-m[j,k])/s; q[i+1]=s/4; q[j+1]=(m[j,i]+m[i,j])/s; q[k+1]=(m[k,i]+m[i,k])/s
        w,x,y,z=q
    return np.array([w,x,y,z])

def ring_colliders(prefix,r0,r1,z0,z1,segments=64):
    out=[]; dr=r1-r0; dz=z1-z0; length=math.hypot(dr,dz)
    rm=(r0+r1)/2
    for i in range(segments):
        a=2*math.pi*(i+.5)/segments; ca,sa=math.cos(a),math.sin(a)
        xx=np.array([-sa,ca,0.]); zz=np.array([dr*ca,dr*sa,dz])/length; yy=np.cross(zz,xx)
        q=mat_to_quat(np.column_stack([xx,yy,zz]))
        p=[rm*ca,rm*sa,(z0+z1)/2]
        tangent=math.pi*rm/segments*1.13
        out.append(f'<geom name="{prefix}_{i:02d}" type="box" size="{tangent:.7f} 0.0008 {length/2+.0003:.7f}" pos="{p[0]:.7f} {p[1]:.7f} {p[2]:.7f}" quat="{q[0]:.8f} {q[1]:.8f} {q[2]:.8f} {q[3]:.8f}" friction="0.16 0.004 0.00005"/>')
    return ''.join(out)
cone_colliders=ring_colliders('wall_seg',THROAT_R,CONE_R1,CONE_Z0,CONE_Z1)
tube_colliders=ring_colliders('throat_seg',THROAT_R,THROAT_R,.0655,.081)
beans=initial_beans()
bean_xml=[]
for i,(p,q) in enumerate(beans,1):
    bean_xml.append(f'''<body name="bean_{i:02d}" pos="{p[0]:.7f} {p[1]:.7f} {p[2]:.7f}" quat="{q[0]:.8f} {q[1]:.8f} {q[2]:.8f} {q[3]:.8f}">
      <freejoint/>
      <geom name="bean_{i:02d}_dome" class="bean" type="ellipsoid" size="0.0030 0.0048 0.0018" pos="0 0 0.00045" mass="0.000115" material="bean_dome"/>
      <geom name="bean_{i:02d}_flat" class="bean" type="ellipsoid" size="0.00285 0.00445 0.00068" pos="0 0 -0.00105" mass="0.000065" material="bean_flat"/>
    </body>''')
xml=f'''<mujoco model="grinder_hopper_50_beans">
  <compiler angle="radian" coordinate="local" inertiafromgeom="true"/>
  <option timestep="{DT}" gravity="0 0 -9.81" integrator="implicitfast" solver="Newton" iterations="100" tolerance="1e-10" cone="elliptic"/>
  <size memory="64M"/>
  <default>
    <geom contype="1" conaffinity="1" solref="0.004 1" solimp="0.90 0.97 0.001" friction="0.65 0.008 0.0001"/>
    <default class="bean"><geom friction="0.32 0.006 0.00008"/></default>
  </default>
  <asset>
    <material name="bean_dome" rgba="0.30 0.105 0.035 1"/>
    <material name="bean_flat" rgba="0.48 0.225 0.085 1"/>
    {cone_mesh}
    {tube_mesh}
  </asset>
  <worldbody>
    {''.join(bean_xml)}
    <geom name="tray" type="box" size="0.074 0.074 0.0015" pos="0 0 -0.0015" material="bean_flat" friction="0.8 0.01 0.0001"/>
    <geom name="rim_front" type="box" size="0.075 0.0012 0.006" pos="0 -0.074 0.0045"/>
    <geom name="rim_back" type="box" size="0.075 0.0012 0.006" pos="0 0.074 0.0045"/>
    <geom name="rim_left" type="box" size="0.0012 0.074 0.006" pos="-0.074 0 0.0045"/>
    <geom name="rim_right" type="box" size="0.0012 0.074 0.006" pos="0.074 0 0.0045"/>
    <geom name="hopper_wall" type="mesh" mesh="hopper_shell" contype="0" conaffinity="0"/>
    <geom name="throat_wall" type="mesh" mesh="throat_shell" contype="0" conaffinity="0"/>
    {cone_colliders}
    {tube_colliders}
    <body name="sliding_gate" pos="0 0 {GATE_Z}">
      <joint name="gate_slide" type="slide" axis="1 0 0" range="0 0.038" damping="0.08"/>
      <geom name="gate_plate" type="box" size="0.017 0.017 0.001" mass="0.003" friction="0.42 0.004 0.00005"/>
    </body>
  </worldbody>
  <actuator><position name="gate_actuator" joint="gate_slide" kp="900" dampratio="1.0" ctrlrange="0 0.038"/></actuator>
</mujoco>'''
XML_PATH.write_text(xml)
model=mujoco.MjModel.from_xml_path(str(XML_PATH))
data=mujoco.MjData(model)

snap_times=[0.0,1.0,2.0,5.0]
snapshots={}
through=[]
passed=np.zeros(N_BEANS,dtype=bool)
passed_counts={}
wall_max=0.0
escaped=set()
diagnostic=None

def capture(t):
    snapshots[t]=dict(qpos=data.qpos.copy(),time=data.time,gate=data.qpos[model.jnt_qposadr[model.joint('gate_slide').id]])
    passed_counts[t]=int(passed.sum())

def contact_record():
    pts=[]; forces=[]; kinds=[]
    ids={mujoco.mjtObj.mjOBJ_GEOM:None}
    for k in range(data.ncon):
        c=data.contact[k]
        f=np.zeros(6)
        mujoco.mj_contactForce(model,data,k,f)
        if np.linalg.norm(f[:3])<1e-7: continue
        pts.append(c.pos.copy()); forces.append(f[:3].copy())
        names={model.geom(c.geom1).name,model.geom(c.geom2).name}
        kinds.append('wall' if names & {'hopper_wall','throat_wall','gate_plate'} else 'other')
    return np.array(pts),np.array(forces),kinds

nsteps=round(END/DT)
for step in range(nsteps+1):
    t=step*DT
    data.ctrl[0]=0.0 if t<1.0 else 0.038
    if any(abs(t-s)<DT/10 for s in snap_times):
        capture(round(t,3))
    if step==nsteps: break
    mujoco.mj_step(model,data)
    for i in range(N_BEANS):
        bid=model.body(f'bean_{i+1:02d}').id
        p=data.xpos[bid]
        if not passed[i] and t>=1.0 and p[2]<COUNT_Z and data.cvel[bid][5]<0:
            passed[i]=True
        if p[2]<-0.025 or abs(p[0])>0.09 or abs(p[1])>0.09:
            escaped.add(f'bean_{i+1:02d}')
    for k in range(data.ncon):
        c=data.contact[k]
        names={model.geom(c.geom1).name,model.geom(c.geom2).name}
        if any(n.startswith(('wall_seg','throat_seg')) for n in names):
            wall_max=max(wall_max,-float(c.dist))
    if abs((t+DT)-2.0)<DT/10:
        diagnostic=contact_record()
    if step%20==0:
        through.append((t,int(passed.sum())))
through.append((END,int(passed.sum())))
if diagnostic is None: diagnostic=contact_record()

print(f"Maximum penetration into hopper walls: {wall_max*1000:.4f} mm")
for t in snap_times:
    print(f"Through throat at {t:.0f} s: {passed_counts[t]} / {N_BEANS}")
print("Beans that left the model: "+(", ".join(sorted(escaped)) if escaped else "none"))
print(f"Final through-throat count: {passed.sum()} / {N_BEANS}; contacts at 2 s: {len(diagnostic[0])}")

# ---------- Drawing helpers ----------
plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'axes.titleweight':'bold'})
BG='#F4F1EA'; INK='#17252B'; MUTED='#617177'; COPPER='#B66A3C'; TEAL='#257A7B'; RED='#D85D43'
fig=plt.figure(figsize=(17,12),facecolor=BG)
gs=GridSpec(2,3,figure=fig,left=.035,right=.98,top=.865,bottom=.105,wspace=.10,hspace=.22,height_ratios=[1,1])
fig.text(.045,.956,'GRINDER HOPPER',fontsize=27,fontweight='bold',color=INK)
fig.text(.045,.918,'50 roasted coffee beans  /  controlled discharge study',fontsize=15,color=COPPER,fontweight='bold')
fig.text(.965,.949,'MUJOCO  •  5.00 s',ha='right',fontsize=12,color=INK,fontweight='bold')
fig.text(.965,.922,'Ø18 mm throat   |   gate opens at 1.00 s',ha='right',fontsize=11,color=MUTED)

# Geometry for drawing.
n=64
ang=np.linspace(0,2*np.pi,n,endpoint=False)
cone_inner=np.array([[THROAT_R*np.cos(a),THROAT_R*np.sin(a),CONE_Z0] for a in ang]+[[CONE_R1*np.cos(a),CONE_R1*np.sin(a),CONE_Z1] for a in ang])
cone_faces=[[cone_inner[i],cone_inner[(i+1)%n],cone_inner[n+(i+1)%n],cone_inner[n+i]] for i in range(n)]
tube_faces=[]
for i in range(n):
    j=(i+1)%n
    tube_faces.append([[THROAT_R*np.cos(ang[i]),THROAT_R*np.sin(ang[i]),.0655],
                       [THROAT_R*np.cos(ang[j]),THROAT_R*np.sin(ang[j]),.0655],
                       [THROAT_R*np.cos(ang[j]),THROAT_R*np.sin(ang[j]),.081],
                       [THROAT_R*np.cos(ang[i]),THROAT_R*np.sin(ang[i]),.081]])
# Low-resolution unit sphere, shared by all ellipsoids.
u=np.linspace(0,2*np.pi,13); v=np.linspace(0,np.pi,8)
unit=np.stack(np.meshgrid(np.cos(u)*np.sin(v)[:,None] if False else u,v),axis=-1) if False else None

def quat_mat(q):
    w,x,y,z=q
    return np.array([[1-2*(y*y+z*z),2*(x*y-z*w),2*(x*z+y*w)],
                     [2*(x*y+z*w),1-2*(x*x+z*z),2*(y*z-x*w)],
                     [2*(x*z-y*w),2*(y*z+x*w),1-2*(x*x+y*y)]])

def draw_ellipsoid(ax,p,q,radii,offset,color,alpha=1):
    uu=np.linspace(0,2*np.pi,13); vv=np.linspace(0,np.pi,8)
    U,V=np.meshgrid(uu,vv)
    xyz=np.stack([radii[0]*np.cos(U)*np.sin(V),radii[1]*np.sin(U)*np.sin(V),radii[2]*np.cos(V)],axis=-1)
    xyz=xyz@quat_mat(q).T + p + quat_mat(q)@np.array(offset)
    ax.plot_surface(xyz[...,0],xyz[...,1],xyz[...,2],color=color,linewidth=0,antialiased=True,shade=True,alpha=alpha)

def base3d(ax,title,subtitle):
    ax.set_facecolor(BG)
    ax.text2D(0,1.09,title,transform=ax.transAxes,fontsize=14,color=INK,fontweight='bold',ha='left')
    ax.text2D(0,1.045,subtitle,transform=ax.transAxes,color=MUTED,fontsize=9,ha='left')
    ax.set_xlim(-.079,.079); ax.set_ylim(-.079,.079); ax.set_zlim(-.002,.19)
    ax.set_box_aspect((1,1,1.12)); ax.view_init(elev=20,azim=-57)
    ax.set_axis_off()

def draw_scene(ax,t,diag=False):
    s=snapshots[t]
    # Tray and raised rim.
    tray=[ [(-.074,-.074,0),(.074,-.074,0),(.074,.074,0),(-.074,.074,0)] ]
    ax.add_collection3d(Poly3DCollection(tray,facecolor='#D4CEC1',edgecolor='#A69E8F',alpha=.16))
    for x in [-.074,.074]:
        ax.plot([x,x],[-.074,.074],[.009,.009],color='#8D958F',lw=2)
    for y in [-.074,.074]: ax.plot([-.074,.074],[y,y],[.009,.009],color='#8D958F',lw=2)
    # Semi-transparent metal shell plus clean silhouette rings.
    if not diag:
        ax.add_collection3d(Poly3DCollection(cone_faces,facecolor='#A7C4BF',edgecolor='#729D97',linewidth=.12,alpha=.17))
    ax.add_collection3d(Poly3DCollection(tube_faces,facecolor='#8AAEAA',edgecolor='#5C8783',linewidth=.15,alpha=.28))
    rings=[(THROAT_R,CONE_Z0),(THROAT_R,.0655)] if diag else [(CONE_R1,CONE_Z1),(THROAT_R,CONE_Z0),(THROAT_R,.0655)]
    for r,z in rings:
        ax.plot(r*np.cos(np.r_[ang,ang[0]]),r*np.sin(np.r_[ang,ang[0]]),np.full(n+1,z),color=TEAL,lw=1.4,alpha=.9)
    # Gate plate (actual simulated joint position).
    gx=s['gate']; gz=GATE_Z
    gate=[[(gx-.017,-.017,gz+.001),(gx+.017,-.017,gz+.001),(gx+.017,.017,gz+.001),(gx-.017,.017,gz+.001)]]
    ax.add_collection3d(Poly3DCollection(gate,facecolor=RED,edgecolor='#923E30',alpha=.92))
    # Bodies: qpos is free-joint position + wxyz quaternion for each bean, then gate.
    for i in range(N_BEANS):
        adr=model.jnt_qposadr[model.body(f'bean_{i+1:02d}').jntadr[0]]
        p=s['qpos'][adr:adr+3]; q=s['qpos'][adr+3:adr+7]
        draw_ellipsoid(ax,p,q,(.003,.0048,.0018),(0,0,.00045),'#82401F')
        draw_ellipsoid(ax,p,q,(.00285,.00445,.00068),(0,0,-.00105),'#B66A35')
    # Throat dimension leader, visible in every view.
    ax.plot([-.009,.009],[0,0],[.064,.064],color=INK,lw=.8)
    ax.text(.012,0,.064,'Ø18',fontsize=8,color=INK)
    if diag:
        pts,forces,kinds=diagnostic
        if len(pts):
            # All physical contacts, with actual normal/tangential resultant vectors.
            ax.scatter(pts[:,0],pts[:,1],pts[:,2],s=5,c='#F4C542',edgecolors='none',depthshade=False)
            # Keep the diagnostic legible: show the 45 largest resultant forces.
            chosen=np.argsort(np.linalg.norm(forces,axis=1))[-45:]
            ax.scatter(pts[chosen,0],pts[chosen,1],pts[chosen,2],s=13,c='#F4C542',edgecolors=INK,linewidths=.25,depthshade=False)
            scale=8.0 # metres per newton; 1 mN -> 8 mm
            for p,f in zip(pts[chosen],forces[chosen]):
                d=f*scale
                if np.linalg.norm(d)>.0003:
                    ax.quiver(*p,*d,color=RED,linewidth=.65,arrow_length_ratio=.25)
        ax.text2D(.02,.04,'● contact    ↗ 45 largest resultant forces  (1 mN = 8 mm)',transform=ax.transAxes,fontsize=8.5,color=INK)

positions=[(0,0),(0,1),(0,2),(1,1)]
for t,(r,c) in zip(snap_times,positions):
    ax=fig.add_subplot(gs[r,c],projection='3d')
    gate_text='CLOSED' if t==0 else ('OPENING' if t==1 else 'OPEN')
    base3d(ax,f'{t:.0f} s',f'{passed_counts[t]:02d} / 50 through throat   •   gate {gate_text}')
    draw_scene(ax,t)
axd=fig.add_subplot(gs[1,0],projection='3d')
base3d(axd,'DIAGNOSTIC  /  2 s',f'{len(diagnostic[0])} active contacts  •  force vectors in world frame')
draw_scene(axd,2.0,diag=True)
axd.set_xlim(-.058,.058); axd.set_ylim(-.058,.058); axd.set_zlim(-.002,.078)
axd.set_box_aspect((1.25,1.25,.72)); axd.view_init(elev=32,azim=-57)

# Throughput panel.
ax=fig.add_subplot(gs[1,2]); ax.set_facecolor(BG)
tv=np.array(through)
ax.step(tv[:,0],tv[:,1],where='post',color=TEAL,lw=2.7)
ax.fill_between(tv[:,0],tv[:,1],step='post',color=TEAL,alpha=.10)
ax.axvspan(0,1,color=RED,alpha=.065)
ax.axvline(1,color=RED,lw=1.3,ls='--')
ax.text(1.06,46,'gate opens',color=RED,fontsize=9,va='top')
for t in snap_times:
    y=passed_counts[t]
    ax.plot(t,y,'o',ms=7,mfc=BG,mec=COPPER,mew=2)
    ax.annotate(f'{y}',(t,y),xytext=(-13 if t==5 else 5,8),textcoords='offset points',fontsize=10,color=INK,fontweight='bold')
ax.set_xlim(0,5); ax.set_ylim(0,55); ax.set_xticks(range(6)); ax.set_yticks([0,10,20,30,40,50])
ax.set_xlabel('time  [s]',color=MUTED); ax.set_ylabel('cumulative beans below throat',color=MUTED)
ax.set_title('BEANS THROUGH THROAT',loc='left',fontsize=14,color=INK,pad=20)
ax.grid(axis='y',color='#C9C7BE',lw=.7,alpha=.7); ax.set_axisbelow(True)
for sp in ax.spines.values(): sp.set_visible(False)
ax.tick_params(colors=MUTED,length=0)

fig.text(.045,.064,'BEAN  9.6 × 6.0 × 4.0 mm  •  0.180 g  •  two overlapping ellipsoids: domed back + flattened face',fontsize=10,color=INK)
fig.text(.045,.040,'FRICTION  bean 0.32 / wall 0.16 / tray 0.80     SOLVER  Newton, elliptic cone, implicitfast, Δt 0.5 ms, 100 iterations',fontsize=9,color=MUTED)
fig.text(.955,.040,f'max wall penetration  {wall_max*1000:.3f} mm    |    escaped beans  {len(escaped)}',fontsize=9,color=MUTED,ha='right')
fig.savefig(OUT,dpi=180,facecolor=BG,bbox_inches='tight')