from QuantumATK import *
from utilities import vectorToGrid, averageFermiLevel
import scipy

# Read the configuration
device_configuration = nlread('au_dtb_au.hdf5', DeviceConfiguration)[0]
# Get H and S
H, S = calculateHamiltonianAndOverlap(device_configuration)
H = H.inUnitsOf(eV)

# Calculate average Fermi level
average_fermi_level = averageFermiLevel(device_configuration)

# Get index of orbitals on the Phenyl ring
projection_list = ProjectionList(elements=[Carbon, Hydrogen])
orbital_index = projection_list.orbitalIndex(device_configuration)
# Project H, S onto the Phenyl ring
H_dtb = H[orbital_index,:][:, orbital_index]
S_dtb = S[orbital_index,:][:, orbital_index]

# Calculate the eigenfunctions
w, v = scipy.linalg.eigh(H_dtb, S_dtb)
# Calculate eigen energies relative to the fermi level
eigen_energies =  (w*eV-average_fermi_level)

# Save eigenstates number 13-17
for i in [8, 11, 15, 17]:
    print('eigenenergy ', i,  eigen_energies[i])
    # Put the eigenvector into a vector of length of all orbitals
    number_orbitals = H.shape[0]
    eigenvector = numpy.zeros(number_orbitals, dtype=complex)
    eigenvector[orbital_index] = v[:,i]
    grid = vectorToGrid(eigenvector, device_configuration)
    nlsave('mpsh.hdf5', grid)