from pprint import pprint
import numpy as np
import scipy.constants as c
from numpy import exp
from math import sin, cos, atan, acos, sqrt, pi
a=c.physical_constants['Bohr radius'][0]
i=complex(0,1)
def complex_round(x,n):
    ##print(type(x))
    reel=x.real
    imaaag=x.imag
    #print(x)
    #print(complex(round(reel,n), round(imaaag,n)))
    return complex(round(reel,n), round(imaaag,n))
def spherical_to_cartesian(r,phi,theta):
    x=r*sin(theta)*cos(phi)
    y=r*sin(phi)*sin(theta)
    z=r*cos(phi)
    if abs(x)< 10e-14: x=0
    if abs(y)< 10e-14: y=0
    if abs(z)< 10e-14: z=0
    #print(x,y,z)
    return x,y,z
    pass
def sci_round(x,n):
    return float(('{:0.'+str(n)+'e}').format(x))
def cartesian_to_spherical(x, y, z):
    r=round(sqrt(x**2+y**2+z**2),4)
    theta=round(atan(y/x),4) if x!=0  else (1.5708*(y/abs(y)) if y!=0 else  0)
    phi=round(acos(z/r),4) if r !=0 else 0
    #print(r,phi,theta)
    return r,phi,theta
    pass


def radial_wave_func(n, l, r):
    a = c.physical_constants['Bohr radius'][0]
    constant = 1
    if n == 1 and l == 0:
        return complex_round(2 * constant * exp(-(r / a)),5)

    elif n == 2 and l == 0:
        return complex_round((1 / sqrt(2)) * constant * (1 - (r / (2 * a))) * exp(-(r / (2 * a))),5)

    elif n == 2 and l == 1:
        return complex_round((1 / sqrt(24)) * constant * (r / a) * exp(-(r / (2 * a))),5)

    elif n == 3 and l == 0:
        return complex_round((2 / (81 * sqrt(3))) * constant * (27 - 18 * (r / a) + 2 * (r / a) ** 2) * exp(-(r / (3 * a))),5)

    elif n == 3 and l == 1:
        return complex_round((8 / (27 * sqrt(6))) * constant * (1 - (r / (6 * a))) * (r / a) * exp(-(r / (3 * a))),5)

    elif n == 3 and l == 2:
        return complex_round((4 / (81 * sqrt(30))) * constant * ((r / a) ** 2) * exp(-(r / (3 * a))),5)

    elif n == 4 and l == 0:
        return complex_round((1 / 4) * constant * (1 - (3 / 4) * (r / a) + (1 / 8) * (r / a) ** 2 - (1 / 192) * (r / a) ** 3) * exp(
            -(r / (4 * a))),5)

    elif n == 4 and l == 1:
        return complex_round((sqrt(5) / (16 * sqrt(3))) * constant * (1 - (1 / 4) * (r / a) + (1 / 80) * (r / a) ** 2) * exp(
            -(r / (4 * a))),5)

    elif n == 4 and l == 2:
        return complex_round((1 / (64 * sqrt(5))) * constant * (r / a) ** 2 * (1 - (1 / 12) * (r / a)) * exp(-(r / (4 * a))),5)

    elif n == 4 and l == 3:
        return complex_round((1 / (768 * sqrt(35))) * constant * (r / a) ** 3 * exp(-(r / (4 * a))),5)

def angular_wave_func(m, l, theta, phi):
    if l == 0 and m == 0:
        return complex_round(sqrt(1 / (4 * pi)),5)
    elif l == 1 and m == 1:
        #print((-sqrt(3 / (8 * pi)) * sin(theta) * exp(i * phi),5))
        return complex_round(-sqrt(3 / (8 * pi)) * sin(theta) * exp(i * phi),5)
    elif l == 1 and m == 0:
        return complex_round(sqrt(3 / (4 * pi)) * cos(theta),5)
        pass
    elif l == 1 and m == -1:
        return complex_round(sqrt(3 / 8 / pi) * sin(theta) * exp(-i * phi),5)
        pass
    elif l == 2 and m == 2:
        return complex_round(sqrt(15 / 32 / pi) * (sin(theta) ** 2) * exp(i * 2 * phi),5)
        pass
    elif l == 2 and m == 1:
        return complex_round((-sqrt(15 / 8 / pi)) * cos(theta) * sin(theta) * exp(i * phi),5)
        pass
    elif l == 2 and m == 0:
        return complex_round(sqrt(5 / 16 / pi) * (3 * (cos(theta) ** 2) - 1),5)
        pass
    elif l == 2 and m == -1:
        return complex_round(sqrt(15 / 8 / pi) * cos(theta) * sin(theta) * exp(-i * phi),5)
        pass
    elif l == 2 and m == -2:
        return complex_round(sqrt(15 / 32 / pi) * (sin(theta) ** 2) * exp(-i * 2 * phi),5)
        pass
    elif l == 3 and m == 3:
        return complex_round(-sqrt(35 / 64 / pi) * (sin(theta) ** 3) * exp(i * 3 * phi),5)
        pass
    elif l == 3 and m == 2:
        return complex_round(sqrt(105 / 32 / pi) * cos(theta) * (sin(theta) ** 2) * exp(i * phi),5)
        pass
    elif l == 3 and m == 1:
        return complex_round(-sqrt(21 / 64 / pi) * sin(theta) * (5 * (cos(theta) ** 2) - 1) * exp(i * phi),5)
        pass
    elif l == 3 and m == 0:
        return complex_round(sqrt(7 / 16 / pi) * (5 * (cos(theta) ** 3) - 3 * cos(theta)),5)
        pass
    elif l == 3 and m == -1:
        return complex_round(sqrt(21 / 64 / pi) * sin(theta) * (5 * (cos(theta) ** 2) - 1) * exp(-i * phi),5)
        pass
    elif l == 3 and m == -2:
        return complex_round(sqrt(105 / 32 / pi) * cos(theta) * (sin(theta) ** 2) * exp(-i * 2 * phi),5)
        pass
    elif l == 3 and m == -3:
        return complex_round(sqrt(35 / 64 / pi) * (sin(theta) ** 3) * exp(-i * 3 * phi),5)
        pass
    else:
        return
    pass
# this function accounts for linear combination of angular wave functions
# (week 2 Lesson 1 slide 27)
def lincomb_angular_wave_func(m, l, theta, phi):
    if m<0:
        return (1j / np.sqrt(2)) * (
                angular_wave_func(m, l, theta, phi) - ((-1) ** float(m)) * angular_wave_func(-m, l, theta, phi))
        pass
    elif m>0:
        return (1 / np.sqrt(2)) * (
                    angular_wave_func(-m, l, theta, phi) + ((-1) **  float(m)) * angular_wave_func(m, l, theta, phi))
        pass
    elif m==0:
        return angular_wave_func(m, l, theta, phi)

def mgrid3d(xstart, xend, xpoints,
            ystart, yend, ypoints,
            zstart, zend, zpoints):
    xr = []
    yr = []
    zr = []
    xval = xstart
    xcount = 0

    # calculate the step size for each axis
    xstep = (abs(xend) + abs(xstart)) / (xpoints - 1)
    ystep = (abs(yend) + abs(ystart)) / (ypoints - 1)
    zstep = 0
    if zpoints != 0 or abs(abs(zend) - abs(zstart)) != 0:
        zstep = abs(abs(zend) - abs(zstart)) / (zpoints - 1)
    if zend < 0 or zstart < 0:
        zstep = abs(abs(zend) + abs(zstart)) / (zpoints - 1)
    #print(zend, zstart, zstep)
    while xcount < xpoints:
        # initialize y points
        yval = ystart
        ycount = 0
        # create temporary variable to store x, y and z list
        xrow = []
        yrow = []
        zrow = []

        while ycount < ypoints:
            # initialize z points
            zval = zstart
            zcount = 0
            # create temporary variable to store the inner x, y, and z list
            inner_xrow = []
            inner_yrow = []
            inner_zrow = []

            while zcount < zpoints:
                # add the points x, y, and z to the inner most list
                inner_xrow.append(xval)
                inner_yrow.append(yval)
                inner_zrow.append(zval)

                # increase z point
                zval += zstep
                zcount += 1
            # add the inner most lists to the second lists
            xrow.append(inner_xrow)
            yrow.append(inner_yrow)
            zrow.append(inner_zrow)

            # increase y point
            yval += ystep
            ycount += 1
        # add the second lists to the returned lists
        xr.append(xrow)
        yr.append(yrow)
        zr.append(zrow)

        # increase x point
        xval += xstep
        xcount += 1
    return (xr, yr, zr)


def mgrid2d(xstart, xend, xpoints, ystart, yend, ypoints):
    # initialize a list to store the grid points that will be returned
    xr = []
    yr = []

    # initialize the first x value
    xval = xstart

    # initialize a variable to count the number of x points
    xcount = 0

    # Calculate the step size for x and y, replace None with the right expression
    xstep = (abs(xend) + abs(xstart)) / (xpoints - 1)
    ystep = (abs(yend) + abs(ystart)) / (ypoints - 1)

    while xcount < xpoints:
        # initialize the first y value, replace None with the right value
        yval = ystart

        # initialize the variable to count the number of y points, replace None with the right value
        ycount = 0

        # initialize temporary lists
        xrow = []
        yrow = []

        while ycount < ypoints:
            # write code to add the current x value to the temporary x list
            xrow.append(xval)

            # write code to add the current y value to the temporary y list
            yrow.append(yval)

            # increase the y value by the step size in y
            yval += ystep

            # increase the counter for the number of y points
            ycount += 1
        # Add the temporary x list to the final x list
        xr.append(xrow)

        # Add the temporary y list to the final y list
        yr.append(yrow)

        # increase x value by the step size in x
        xval += xstep

        # increase the counter for the number of x points
        xcount += 1

    return xr, yr

def hydrogen_wave_func(n,l,m,roa,Nx,Ny,Nz):
    #vectorizing functions to be used on arrays
    vec_cart_sphere = np.vectorize(cartesian_to_spherical)
    vec_radial_wave_func = np.vectorize(radial_wave_func)
    vec_linearcombi_angular_wave_func = np.vectorize(lincomb_angular_wave_func)
    array_complex_round = np.vectorize(complex_round)
    array_round = np.vectorize(round)

    xstep = roa*2/(Nx-1)
    xrange = [round(-roa+xstep*i,5) for i in range(Nx)]
    ystep = roa*2/(Ny-1)
    yrange = [round(-roa+ystep*i ,5)for i in range(Ny)]
    zstep = roa*2/(Nz-1)
    zrange = [round(-roa+zstep*i,5) for i in range(Nz)]

    xx = [[[x for z in zrange] for y in yrange] for x in xrange]
    yy = [[[y for z in zrange] for y in yrange] for x in xrange]
    zz = [[[z for z in zrange] for y in yrange] for x in xrange]

    #convert array to spherical
    r,theta,phi = vec_cart_sphere(xx,yy,zz)

    a = c.physical_constants['Bohr radius'][0]

    #get radial function output
    radial = vec_radial_wave_func(n,l,r*a)

    #get angular function output, linearly combined
    angular = vec_linearcombi_angular_wave_func(m,l,theta,phi)

    #get the square of the magnitude of the wavefunction, rounded to 5
    mag = array_round(abs(np.multiply(radial, angular))**2,5)

    #convert nparray back to python list
    return np.round((xx,yy,zz,mag.tolist()),5)