import numpy as np
from mayavi import mlab

def binomial(n, k):
    result = 1
    for i in range(1, k+1):
        result = result * (n-i+1) / float(i)
    return result


def rose(n,d,p):
	if (n*d)%2==1:
		k=1
	else:
		k=2
	d=float(d)
	fi,theta = np.mgrid[0:k*d*np.pi:120j, 0:2*np.pi:120j]
	x=1/2.0*np.cos(fi)*(np.cos(n/d*fi)+np.sqrt(p**2+np.cos(n/d*fi)**2)*np.sin(theta))
	y=1/2.0*np.sin(fi)*(np.cos(n/d*fi)+np.sqrt(p**2+np.cos(n/d*fi)**2)*np.sin(theta))
	z=1/2.0*(p+np.sqrt(p**2+np.cos(n/d*fi)**2)*np.cos(theta))
	return [x,y,z]

	
def implicitni_rose(n,d,p,x0,x1,y0,y1,z0,z1):
    if (n*d)%2==1:
        s=1
    else:
        s=2
    x,y,z = np.mgrid[x0:x1:120j, y0:y1:120j, z0:z1:120j]
    lijevo=0*x
    for k in xrange(d/2+1):
        for j in xrange(0,k+1):
            lijevo += (-1)**(k+j)*binomial(d,2*k)*binomial(k,j)*(x**2+y**2+z**2-p*z)**(d-2*(k-j))*(x**2+y**2)**(k-j)
    lijevo = lijevo**s
    desno = 0*y
    for i in xrange(n/2+1):
        desno += (-1)**i*binomial(n,2*i)*x**(n-2*i)*y**(2*i)
    desno = desno**s
    if n>d:
         rezultat=(x**2+y**2)**(s*(n-d)/2.0)*lijevo-desno
    elif n<d:
         rezultat=lijevo-(x**2+y**2)**(s*(d-n)/2.0)*desno
    return rezultat

#parametarski
#tocke = rose(7,3,0)
#tocke = rose(3,7,0)
#tocke=rose(5,2,1)
#mlab.mesh(tocke[0], tocke[1], tocke[2], colormap='autumn')

#implicitno
#tocke=implicitni_rose(7,3,0,-0.57,0.57,-0.57,0.57,-0.5,0.2)
tocke=implicitni_rose(5,2,1,-1.5,1.5,-1.5,1.5,-0.5,0.2)
#tocke=implicitni_rose(3,7,0,-1,1,-1,1,-0.5,0.2)
mlab.contour3d(tocke, contours=[0], colormap='YlGnBu')


mlab.show()