import pylab

from experimental_tamr import data

filename_out_of_plane = 'W8-Fe2-Co-2x2-theta_0-phi_0.hdf5'
filename_in_plane = 'W8-Fe2-Co-2x2-theta_90-phi_150.hdf5'

ldos_out_of_plane = nlread(filename_out_of_plane, LocalDensityOfStates)[0]
ldos_in_plane = nlread(filename_in_plane, LocalDensityOfStates)[0]


conf = ldos_in_plane._configuration()
R0 = conf.cartesianCoordinates()[-1, :]

z_shift = 2 * Ang

R = R0.copy()
R[2] += z_shift

nlprint('Tip position')
nlprint(R)

n_out = ldos_out_of_plane.evaluate(x=R[0], y=R[1], z=R[2], spin=Spin.Sum)
n_in = ldos_in_plane.evaluate(x=R[0], y=R[1], z=R[2], spin=Spin.Sum)


energies = ldos_in_plane.energies().inUnitsOf(eV)

# Calculate the TAMR
tamr_ldos = (n_out - n_in) / n_out * 100

# Plot the LDOS itself
pylab.figure()
pylab.subplot(211)
pylab.plot(energies, n_in, 'b-', label='$\\parallel$')
pylab.plot(energies, n_out, 'r--', label='$\\perp$')
pylab.plot([0, 0], [0, 1.1*max(n_out)], 'k--')
pylab.ylabel('LDOS')
pylab.legend(loc=0)
pylab.xlim([-0.2, 0.5])

pylab.subplot(212)
pylab.plot(energies, tamr_ldos, 'k-')
pylab.plot([-0.2, 1], [0, 0], 'k--')
pylab.plot([0, 0], [-30, 20], 'k--')
pylab.xlabel('Energy  (eV)')
pylab.ylabel('TAMR (%)')
pylab.xlim([-0.2, 0.5])

pylab.show()