import numpy
import scipy.special
import scipy.misc
import npy2cube

def w(n,l,m,d,is_abs=False):

    x,y,z = numpy.mgrid[-d:d:30j,-d:d:30j,-d:d:30j]
    
    r = lambda x,y,z: numpy.sqrt(x**2+y**2+z**2) # Евклидово расстояние
    theta = lambda x,y,z: numpy.arccos(z/r(x,y,z)) # Вычисление полярного угла
    phi = lambda x,y,z: numpy.arctan(y/x) # Вычисление азимута

    a0 = 1.
    
    R = lambda r,n,l: (2*r/n/a0)**l * numpy.exp(-r/n/a0) * scipy.special.genlaguerre(n-l-1,2*l+1)(2*r/n/a0) # Радиальная часть волновой функции
    WF = lambda r,theta,phi,n,l,m: R(r,n,l) * scipy.special.sph_harm(m,l,phi,theta) # Сама волновая функция
    absWF = lambda r,theta,phi,n,l,m: numpy.absolute(WF(r,theta,phi,n,l,m))**2 # Квадрат волновой функции
    
    return WF(r(x,y,z),theta(x,y,z),phi(x,y,z),n,l,m)      
    

import __main__
__main__.pymol_argv = [ 'pymol', '-x' ]

import pymol
from pymol import cmd
pymol.finish_launching()

import time
from IPython.display import Image

d,step = 15,30

for n in range(0,4):
    for l in range(0,n):
        for m in range(0,l+1,1):
            print(1)
            grid = w(n=n,l=l,m=m,d=d)             
            name ='%s-%s-%s' % (n,l,m)

#             ipv.pylab.figure(debug=True)
#             ipv.pylab.volshow(w(n,l,m,d,True),lighting=True,downscale=4)
#             ipv.pylab.show()
#             ERROR

            npy2cube.npy2cube(grid,(-d,-d,-d),(step,step,step),name+'.cube')

cmd.set("ray_volume")

### cut below here and paste into script ###
cmd.volume_ramp_new('ramp007', [\
      0.005, 0.00, 0.00, 1.00, 0.00, \
      0.01,  0.07, 1.00, 1.00, 0.20, \
      0.015, 1.00, 0.00, 1.00, 1.00, \
    ])
### cut above here and paste into script ###

for n in range(0,4):
    for l in range(0,n):
        for m in range(0,l+1,1):
            name='%s-%s-%s' % (n,l,m)
            
            cmd.load(name+'.cube')
            cmd.volume('vol', name)
            cmd.volume_color('vol','ramp007')
            time.sleep(1)
            cmd.ray()
            time.sleep(1)
            cmd.png(name+'.png')
            time.sleep(1)
            
            display(Image(name+'.png'))

%%bash
jupyter nbconvert --to html Untitled.ipynb