"""Mobile home robot v2: DIFFERENTIAL drive (2 driven wheels mid + 2 casters) + vertical lift mast + 4-DOF
arm reaching forward. Y-up model, world joint -> Z-up, wheels on ground."""
import xml.etree.ElementTree as ET
R=ET.Element("robot",{"name":"mobile_base"})
def link(n): return ET.SubElement(R,"link",{"name":n})
def mat(v,rgba,k): ET.SubElement(ET.SubElement(v,"material",{"name":"m"+k}),"color",{"rgba":rgba})
def bvis(lk,size,xyz,rgba,rpy="0 0 0"):
    v=ET.SubElement(lk,"visual");ET.SubElement(v,"origin",{"xyz":xyz,"rpy":rpy});ET.SubElement(ET.SubElement(v,"geometry"),"box",{"size":size});mat(v,rgba,size+xyz)
def cvis(lk,r,l,xyz,rgba,rpy):
    v=ET.SubElement(lk,"visual");ET.SubElement(v,"origin",{"xyz":xyz,"rpy":rpy});ET.SubElement(ET.SubElement(v,"geometry"),"cylinder",{"radius":r,"length":l});mat(v,rgba,r+xyz)
def svis(lk,r,xyz,rgba):
    v=ET.SubElement(lk,"visual");ET.SubElement(v,"origin",{"xyz":xyz,"rpy":"0 0 0"});ET.SubElement(ET.SubElement(v,"geometry"),"sphere",{"radius":r});mat(v,rgba,r+xyz)
def mvis(lk,fn,xyz,rgba,rpy):
    v=ET.SubElement(lk,"visual");ET.SubElement(v,"origin",{"xyz":xyz,"rpy":rpy});ET.SubElement(ET.SubElement(v,"geometry"),"mesh",{"filename":fn,"scale":"0.001 0.001 0.001"});mat(v,rgba,fn+xyz)
def joint(n,typ,par,ch,xyz,axis=None,lo=None,hi=None,rpy="0 0 0"):
    j=ET.SubElement(R,"joint",{"name":n,"type":typ});ET.SubElement(j,"parent",{"link":par});ET.SubElement(j,"child",{"link":ch});ET.SubElement(j,"origin",{"xyz":xyz,"rpy":rpy})
    if axis:ET.SubElement(j,"axis",{"xyz":axis})
    if lo is not None:ET.SubElement(j,"limit",{"lower":str(lo),"upper":str(hi),"effort":"300","velocity":"1"})
CH="0.30 0.32 0.36 1"; WH="0.1 0.1 0.11 1"; CS="0.2 0.2 0.22 1"; EX="0.62 0.62 0.66 1"; DK="0.26 0.26 0.30 1"; SV="0.34 0.34 0.40 1"; TL="0.9 0.52 0.16 1"
WR=0.06
link("world"); joint("world_up","fixed","world","base_link","0 0 0",rpy="1.5708 0 0")
base=link("base_link")
bvis(base,"0.40 0.10 0.30","0 0.10 0",CH); bvis(base,"0.04 0.85 0.04","0 0.55 -0.13",EX)
# 2 DRIVEN wheels (mid, big) about X
for i,x in enumerate([0.21,-0.21]):
    w=link(f"drive{i}"); cvis(w,str(WR),"0.035","0 0 0",WH,"0 1.5708 0"); joint(f"j_drive{i}","continuous","base_link",f"drive{i}",f"{x} {WR} 0",axis="1 0 0")
# 2 casters (front/back), sphere on a small swivel post
for i,z in enumerate([0.13,-0.13]):
    c=link(f"caster{i}"); svis(c,"0.03",f"0 0 0",CS); joint(f"j_caster{i}","fixed","base_link",f"caster{i}",f"0 0.03 {z}")
# lift + arm (as v1)
lc=link("lift"); bvis(lc,"0.10 0.06 0.10","0 0 0",DK); bvis(lc,"0.02 0.10 0.10","-0.06 0 0",EX)
joint("lift","prismatic","base_link","lift","0.03 0.25 -0.11",axis="0 1 0",lo=0.0,hi=0.65)
prev="lift"
for i,(nm,ax,dz) in enumerate([("marm1","0 1 0",0.06),("marm2","1 0 0",0.055),("marm3","1 0 0",0.055),("marm4","0 0 1",0.05)]):
    lk=link(nm); mvis(lk,"st3215.stl","0 0 0.015",SV,"0 0 0"); joint(f"j_{nm}","revolute",prev,nm,f"0 0 {dz}" if i>0 else "0.02 0 0.06",axis=ax,lo=-1.57,hi=1.57); prev=nm
tool=link("mtool"); bvis(tool,"0.03 0.03 0.05","0 0 0.03",TL); joint("j_mtool","fixed",prev,"mtool","0 0 0.05")
ET.indent(ET.ElementTree(R),"  "); ET.ElementTree(R).write("mobile_base.urdf",encoding="unicode",xml_declaration=True)
print("wrote mobile_base.urdf (differential drive)")