#!/usr/bin/env python
# -*- coding: utf-8 -*- 

import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np

from mpl_toolkits.mplot3d import Axes3D

"""
In 1926, Erwin Schrodinger advanced the famous wave equation that relates the energy of a system to its wave properties.
Because its application to the hydrogen atom is rather complicated, we shall first use wave equation to solve the particle-in-a-box.
The Schrodinger Wave equation expressing in 1D is

d^2ψ/dx^2 + 8π^2m/h^2 (E - V)ψ = 0
"""

# Constants
h = 6.626e-34
m = 9.11e-31

# Values for L and x
x_list = np.linspace(0,1,100)
L = 1

def psi(n,L,x):
    """
    Calculate psi of the wavefunction.
    """
    return np.sqrt(2/L)*np.sin(n*np.pi*x/L)

def psi_2(n,L,x):
    """
    Calculate psi squared of the wavefunction.
    """
    return np.square(psi(n,L,x))

plt.figure(figsize=(15,10))
plt.suptitle("Wave Functions", fontsize=18)

for n in range(1,4):
    # Empty lists for energy and psi wave
    psi_2_list = []
    psi_list = []
    for x in x_list:
        psi_2_list.append(psi_2(n,L,x))
        psi_list.append(psi(n,L,x))
    plt.subplot(3,2,2*n-1)
    plt.plot(x_list, psi_list, color="green")
    plt.xlabel("L", fontsize=13)
    plt.ylabel("psi", fontsize=13)
    plt.xticks(np.arange(0, 1, step=0.5))
    plt.title("n="+str(n), fontsize=16)
    plt.grid()
    plt.subplot(3,2,2*n)
    plt.plot(x_list, psi_2_list, color="red")
    plt.xlabel(u"L", fontsize=13)
    plt.ylabel("psi*psi", fontsize=13)
    plt.xticks(np.arange(0, 1, step=0.5))
    plt.title("n="+str(n), fontsize=16)
    plt.grid()
plt.tight_layout(rect=[0, 0.03, 1, 0.95])

"""
1s orbital
"""

def prob_1s(x,y,z):
    """
    Probability we get a 1. 
    """
    r=np.sqrt(np.square(x)+np.square(y)+np.square(z))
    #Remember.. probability is psi squared!
    return np.square(np.exp(-r)/np.sqrt(np.pi))

# Random coordinates
x=np.linspace(0,1,30)
y=np.linspace(0,1,30)
z=np.linspace(0,1,30)
elements = []
probability = []
for ix in x:
    for iy in y:
        for iz in z:
            #Serialize into 1D object
            elements.append(str((ix,iy,iz)))
            probability.append(prob_1s(ix,iy,iz))

# Ensure sum of probability is 1
probability = probability/sum(probability)

# Getting electron coordinates based on probabiliy
coord = np.random.choice(elements, size=100000, replace=True, p=probability)
elem_mat = [i.split(",") for i in coord]
elem_mat = np.matrix(elem_mat)
x_coords = [float(i.item()[1:]) for i in elem_mat[:,0]]
y_coords = [float(i.item()) for i in elem_mat[:,1]]
z_coords = [float(i.item()[0:-1]) for i in elem_mat[:,2]]

# Plotting
fig = plt.figure(figsize=(10,10))
ax = Axes3D(fig)
ax = fig.add_subplot(111, projection="3d")
ax.scatter(x_coords, y_coords, z_coords, alpha=0.05, s=2)
ax.set_title("Hydrogen 1s density")
plt.show()
