%load_ext autoreload
%autoreload 2

import numpy as np
from numpy.linalg import det
from numpy.linalg import inv
from numpy.fft import fftn, ifftn
from itertools import product
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from pydft.crystal import *
from pydft.bases.planewave import *

sigma1 = 0.75
sigma2 = 0.50

lat_const = 6
r = [lat_const,0,0]
S = [16,16,16]
R = makeR(r)
origin = 1./2*np.sum(R,1)
crystal = Crystal(R,S)
g(crystal)

# Create vector of basis expansion coefficients for the electron density.
n = -gauss_charge_dist(crystal._r,origin,sigma1) + (
    gauss_charge_dist(crystal._r,origin,sigma2))

"""
Explanation: Poisson Solver
This notebook demonstrates how the module crystal and planewave can be used in practice. It creates a convergence plot of the the total energy vs. sampling density, and plots the charge density.
End of explanation
"""


# Plane refers to one of the pages of the n and r matrices.
plane=0

npl = n[plane*S[1]*S[2]:(plane+1)*S[1]*S[2]]
rpl = crystal._r[plane*S[1]*S[2]:(plane+1)*S[1]*S[2]]
xpl = [rpl[i,0] for i in range(len(rpl))]
ypl = [rpl[i,1] for i in range(len(rpl))]

fig = plt.figure()
ax = plt.gca(projection='3d')
ax.scatter(xpl,ypl,list(npl))
plt.title("charge density distribution")
plt.show()
plt.close()

ϕ=-4*np.pi*I(crystal, Linv(crystal, O(crystal, J(crystal, n))))
UNumeric = 1./2*np.real(np.dot(J(crystal, ϕ),O(crystal, J(crystal, n))))
UNumeric

# Define the Gaussian broadening parameters
sigma1 = 0.75
sigma2 = 0.50

# Find the analytic solution.
U_analytic = ((1/sigma1 + 1/sigma2)/2 - np.sqrt(2)/
                np.sqrt(sigma1**2 + sigma2**2))/np.sqrt(np.pi)
Uerrors = []
ss = range(8,15)
for si in ss:
    # Create integer, sampling matrices.
    S = [si,si,si]

    # Define Coordinates.
    lat_const = 6
    r = [lat_const,0,0]
    R = makeR(r)
    origin = 1./2*np.sum(R,1)
    
    # Create the crystal
    crystal = Crystal(R, S)
    g(crystal)
    
    # Create vector of basis expansion coefficients for the electron density.
    n = -gauss_charge_dist(crystal._r, origin,sigma1) + (
        gauss_charge_dist(crystal._r,origin,sigma2))

    ϕ = -4*np.pi*I(crystal, Linv(crystal, O(crystal, J(crystal, n))))
    U_numeric = 1./2*np.real(np.dot(J(crystal, ϕ),O(crystal, J(crystal, n))))


    Uerrors.append(abs(U_numeric-U_analytic))
    del crystal
plt.semilogy(ss, Uerrors)
plt.xlabel("Sampling points per edge")
plt.ylabel("Error")
plt.show()
plt.close()

"""
Explanation: Plot the charge density on a plane lying along the x-y axis.
End of explanation
"""


from pydft.crystal import *
from pydft.energies.ewald import *

import numpy as np

# We should get something close to what Arias got with similar parameters.
LC = 6
r = [LC,0,0]
R = makeR(r)
S = [15,15,15]
crystal = Crystal(R, S, LC)
crystal._X = [[0., 0., 0.]]
energy = ewald_energy(crystal)
print(energy)

"""
Explanation: Ewald Sum
End of explanation
"""


error = abs(energy - (-1./3))
print(error)

"""
Explanation: In the assignment Arias got -1./3 so this is fairly close
End of explanation
"""


from pydft.energies.schrodinger import *
from pydft.crystal import *

LC = 1.
R = makeR([LC, 0., 0.])
S = [2,2,2]
crystal = Crystal(R, S, LC)
basis = "planewave"
potential = "sho"
w = 1.
E = getE(crystal, basis, potential, [w])

E

"""
Explanation: Energy from Solving Schrodinger
End of explanation
"""


