import os
import psi4
import matplotlib.pyplot as plt
import numpy as np
# import numpy_html
psi4.set_options({"save_jk" : True})
psi4.set_memory(int(2.50e9))
psi4.core.clean()

import n2v

import matplotlib as mpl
mpl.rcParams["font.size"] = 11
mpl.rcParams["font.family"] = "sans-serif"
mpl.rcParams["axes.edgecolor"] = "#eae8e9" 

#Define Psi4 geometries. Symmetries need to be set to C1.


Ne = psi4.geometry(
"""
0 1
Ne 0.0 0.0 0.0
noreorient
nocom
units bohr
symmetry c1
""" )

#n2v is driven by psi4's reference option. Make sure you set it accordingly.
psi4.set_options({"reference" : "rhf"})

#Perform a calculation for a target density.
#Remember that for post scf calculations, Psi4 does not update the density.
#Thus make sure you obtain something like a dipole in order to do so.
e, wfn = psi4.properties("CCSD/aug-cc-pvtz", return_wfn=True, properties=["dipole"], molecule=Ne)

arrDa = wfn.Da().to_array()
arrDb = wfn.Db().to_array()

diff = arrDa - arrDb
sum_difference = sum(sum(row) for row in diff)
 
#Define inverter objects for each molcule. Simply use the wnf object from psi4 as an argument.
inv = n2v.Inverter('psi4')
inv.set_system(Ne, 'aug-cc-pvtz', wfn=wfn)
inv.Dt = [ np.array(wfn.Da()), np.array(wfn.Db()) ]
inv.ct = [ np.array(wfn.Ca_subset("AO", "OCC")), np.array(wfn.Cb_subset("AO", "OCC")) ]
inv.et = [ np.array(wfn.epsilon_a_subset("AO", "OCC")), np.array(wfn.epsilon_b_subset("AO", "OCC"))]

# Additionally one can simply initialize an Inverter using the wavefunction.
inv = n2v.Inverter.from_wfn(wfn)

# Let us define a plotting grid:

npoints=1001
x = np.linspace(-5,5,npoints)[:,None]
y = np.zeros_like(x)
z = y
grid = np.concatenate((x,y,z), axis=1).T

mix = [0.0, 0.1, 0.5, 1.0]
vxc_lab = ['Vxc_mix0', 'Vxc_mix0.5', 'Vxc_mix_1.0']
vxc_dic = {}
for m in mix:
    inv.invert("zmp", opt_max_iter=200, opt_tol=1e-7, zmp_mixing=m,
           lambda_list=np.linspace(10, 1000, 20), guide_components="fermi_amaldi")
    inv.eigvecs_a[:inv.nalpha]
    np.diag(inv.Da)[:inv.nalpha]
    results = inv.eng.grid.esp(Da=inv.proto_density_a, Db=inv.proto_density_b, grid=grid, )
    vxc_dic[m] = results[1]
    


for k,v in vxc_dic.items():
    plt.plot(x, v, label="Vxc_mix_"+str(k))

plt.legend()
plt.xlim(-5,5)

fig, ax = plt.subplots()
ls = ["solid","--", "-.", "-."]
i = 0
for k,v in vxc_dic.items():
    ax.plot(x,  v, label="vxc_mix_"+str(k), ls=ls[i])
    i += 1

ax.set_title("Neon Exchange Correlation Potenial")
ax.legend()
ax.set_xlim(1e-5,5)
ax.set_xscale("log")



#ax.set_title("Neon Exchange Correlation Potential")

#ax.legend()
#ax.set_xlim(-5,5)

pltname = '_Vxc_'  + '.pdf'
plt.savefig(pltname)
plt.close()


plt.show()
plt.close()