import numpy
import scipy.special
import scipy.misc
import npy2cube
import IPython
import math
import matplotlib.pyplot as plt
from xmlrpc.client import ServerProxy
from IPython.display import Image
import os, sys, time

def w(n, l, m, d):
    
    x, y, z = numpy.mgrid[-d:d:30j, -d:d:30j, -d:d:30j]

    r = numpy.sqrt(x**2 + y**2 + z**2)
    theta = numpy.arccos(z / r)
    phi = numpy.arctan2(y, x)

    a0 = 1.

    norm_coef = math.sqrt((2/n/a0)**3 * math.factorial(n-l-1)/(2*n*math.factorial(n+l)))

    R = (2*r/n/a0)**l * numpy.exp(-r/n/a0) * scipy.special.genlaguerre(n-l-1, 2*l+1)(2*r/n/a0)

  
    WF = norm_coef * R * scipy.special.sph_harm(m, l, phi, theta)

    absWF = numpy.absolute(WF)**2

    return numpy.real(WF) + numpy.imag(WF)

where_am_i = os.getcwd()

for n in range(1,4):
    d = 6.5*n
    step = 2*d/30
    for l in range(0,n):
        for m in range(-l,l+1):
            grid = w(n,l,m,d)
            name = '{}_{}_{}'.format(n, l, m)
            file_path = os.path.join(where_am_i, 'wf_cube', name + '.cube')
            npy2cube.npy2cube(grid, (-d,-d,-d), (step,step,step), file_path)

import __main__

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

import pymol 
pymol.finish_launching()

from pymol import cmd,stored

cmd = ServerProxy(uri="http://localhost:9123/RPC2")

twist_90 = ['3_1_0', '3_1_1', '3_1_-1', '3_2_2', '3_2_-2']

cmd.do("cmd.reinitialize()")
cmd.do("cmd.bg_color('white')")

cmd.volume_ramp_new('my_ramp', [-0.015, 1.00, 1.00, 0.00, 0.50, 
                                -0.01,  0.20, 0.80, 0.00, 0.20, 
                                -0.005, 0.60, 0.60, 0.00, 0.00, 
                                0.005, 0.80, 0.40, 0.00, 0.00, 
                                0.01,  0.80, 0.20, 1.00, 0.20, 
                                0.015, 1.00, 0.00, 1.00, 0.50])
    
for n in range(1,4):
    for l in range(0,n):
        for m in range(-l,l+1):
            name='{}_{}_{}'.format(n, l, m)
            fname = 'wf_cube/' + name + '.cube'
            vname = name + '_vol'
            cmd.do("cmd.load(fname, name)")
            cmd.do("cmd.volume(vname, name, ramp='my_ramp')")
            if twist_90[m] in vname:
                cmd.do("cmd.rotate('y', 90)")
            cmd.do("cmd.hide('all')")
            cmd.do("cmd.show('volume', vname)")
            cmd.do("cmd.draw(720, 720, antialias=2)")
            cmd.do("cmd.png('wf_cube/pics/' + name)")

f, ax = plt.subplots(5, 3, figsize=(10, 10))

row, col = 0, 0
for n in range(1, 4):
    for l in range(0, n):
        for m in range(-l, l+1):
            name = '{}_{}_{}'.format(n, l, m)
            pic_path = os.path.join('wf_cube', 'pics', name)
            f = plt.imread('{}.png'.format(pic_path))
            ax[row, col].imshow(f)
            ax[row, col].axis('off')
            ax[row, col].set_title('{}'.format(name))
            col += 1
            if col == 3: row += 1; col = 0

ax[4,2].set_visible(False)    
plt.tight_layout()

with open('orca_file.inp', 'w') as orca_file:
    orca_file.write('''! UHF SVP XYZFile
%plots Format Cube
MO("H_1.cube",1,0);
MO("H_2.cube",2,0);
MO("H_3.cube",3,0);
MO("H_4.cube",4,0);
MO("H_5.cube",1,1);
MO("H_6.cube",2,1);
MO("H_7.cube",3,1);
MO("H_8.cube",4,1);
end
* xyz 0 2
H 0 0 0
*''')

import subprocess

subprocess.run('echo export PATH=$HOME/orca:$PATH >> ~/.bash_profile', shell=True)
subprocess.run('. ~/.bash_profile', shell=True)
subprocess.run('orca orca_file.inp > orca_file.out', shell=True)

cmd.do("cmd.reinitialize()")
cmd.do("cmd.bg_color('white')")

cmd.volume_ramp_new('my_ramp', [-0.015, 1.00, 1.00, 0.00, 0.50, 
                                -0.01,  0.20, 0.80, 0.00, 0.20, 
                                -0.005, 0.60, 0.60, 0.00, 0.00, 
                                0.005, 0.80, 0.40, 0.00, 0.00, 
                                0.01,  0.80, 0.20, 1.00, 0.20, 
                                0.015, 1.00, 0.00, 1.00, 0.50])
    
for n in range(1,9):
    name='H_{}'.format(n)
    fname = 'orca_cube/' + name + '.cube'
    vname = name + '_vol'
    cmd.do("cmd.load(fname, name)")
    cmd.do("cmd.volume(vname, name, ramp='my_ramp')")
    cmd.do("cmd.hide('all')")
    cmd.do("cmd.show('volume', vname)")
    cmd.do("cmd.draw(720, 720, antialias=2)")
    cmd.do("cmd.png('orca_cube/pics/' + name)")

f, ax = plt.subplots(2, 4, figsize=(10, 3))

row, col = 0, 0
for i in range(1,9):
    name = 'H_{}'.format(i)
    pic_path = os.path.join('orca_cube', 'pics', name)
    f = plt.imread('{}.png'.format(pic_path))
    ax[row, col].imshow(f)
    ax[row, col].axis('off')
    ax[row, col].set_title('{}'.format(name))
    col += 1
    if col == 4: row += 1; col = 0

plt.tight_layout()