# Copyright Tony Kaap 2008
# ballStick.py
# This can be used for any legal purpose, as long as this attribution stays
# intact.

# example usages in the Maya Python Script Editor

#import ballStick
#reload(ballStick)
#ballStick.ballStick()

#import ballStick
#reload(ballStick)
#ballStick.ballStick("pSolidShape1")


import maya.cmds as m
import re
from math import *

point_distance = lambda a,b: ((a[0]-b[0])*(a[0]-b[0])+(a[1]-b[1])*(a[1]-b[1]))

def vector(v1,v2):
	return [v2[0]-v1[0],v2[1]-v1[1],v2[2]-v1[2]]
	
def magnitude(v):
	m2 = 0
	for i in xrange(len(v)):
		  m2 += (v[i]*v[i])
	return sqrt(m2)

def midpoint(pt0,pt1):
	return( [(pt0[0]+pt1[0])/2.0, (pt0[1]+pt1[1])/2.0, (pt0[2]+pt1[2])/2.0] )
	
def ballStickSelected(r=0.05,secX=10,secY=5):
	if not obj:
		obj = m.ls(sl=1)[0]
	ballStick(obj,r,secX,secY)
	
def ballStick(obj=None,r=0.05,secX=10,secY=5):
	r_2 = r/2.0
	if not obj:
		obj = m.ls(sl=1)[0]
	
	print "using: " + obj
	v = m.polyEvaluate(obj,v=1)
	grp = m.group(em=1, name="ballGroup1")
	m.addAttr(grp, ln='radius', at="float", min=0 )
	m.setAttr(grp+'.radius',r,e=1,keyable=1)
	m.addAttr(grp, ln="secX", at="long", min=0 )
	m.setAttr(grp+".secX",secX,e=1,keyable=1)
	m.addAttr(grp, ln="secY", at="long", min=0 )
	m.setAttr(grp+".secY",secY,e=1,keyable=1)
	
	cGrp = m.group(em=1, name="stickGroup1")
	m.addAttr(cGrp, ln='radius', at="float", min=0 )
	m.setAttr(cGrp+'.radius',r_2,e=1,keyable=1)
	m.addAttr(cGrp, ln="sec", at="long", min=0 )
	m.setAttr(cGrp+".sec",secX,e=1,keyable=1)

	pts=[]
	for i in xrange(v):
		pts.append(m.xform(obj+".vtx["+str(i)+"]", query=1, translation=1, ws=1))
		
	#add spheres
	
	for i in xrange(v):
		pos = pts[i]
		(sph,sphNode) = m.polySphere(radius=r, sx=secX,sy=secY)
		m.connectAttr(grp+".radius",sphNode+".radius")
		m.connectAttr(grp+".secX",sphNode+".subdivisionsAxis")
		m.connectAttr(grp+".secY",sphNode+".subdivisionsHeight")

		m.parent(sph,grp)
		m.xform(sph, ws=1, translation=pos)
	
	# add cylinders
	e = m.polyEvaluate(obj,e=1)
	space = re.compile(' +')
	for i in xrange(e):
		v2 = m.polyInfo(obj+".e["+str(i)+"]", edgeToVertex=1)
		v2 = v2[0]
		v2 = space.sub(',',v2) # replace gaps of whitespace with a single comma
		v2 = v2.split(',')
		p0 = int(v2[2])
		p1 = int(v2[3])
		pt0 = pts[p0]
		pt1 = pts[p1]
		mpt = midpoint(pt0,pt1)
		mag = magnitude(vector(pt0,pt1))
		(cyl,cylNode) = m.polyCylinder(height=mag,sy=0,sx=secX,radius=r_2,ax=(0,0,1))
		m.connectAttr(cGrp+".radius",cylNode+".radius")
		m.connectAttr(cGrp+".sec",cylNode+".subdivisionsAxis")
		#h_2 = mag/2.0
		
		m.xform(cyl, ws=1,translation=mpt)
		m.parent(cyl,cGrp)
		# determine the orientation of a z-axis cylinder given two world-space 
		# input points
		rot = orient(pt0,pt1)
		m.rotate(rot[0], rot[1], rot[2] , cyl ,a=1)
		m.select([grp,cGrp])
		
		
def normalize(v):
	ret = v
	mag = magnitude(v)
	for i in range(len(ret)):
		ret[i] = ret[i]/mag
	return ret

def orient(pt0,pt1):
	radToDeg = 180/pi
	d = normalize(vector(pt0,pt1))
	[dx,dy,dz] = d	
	rx=ry=rz=0
	
	# phi - the angle of elevation
	phi = radToDeg*asin(dy)
	if(dz>0):
		phi*=-1
	
	#theta 	- the angle around the y axis
	if dz==0:
		if dx > 0:
			theta = -90
		else:
			theta = 90
	else:
		theta = radToDeg*atan(dx/dz)
	
	return (phi,theta,0)