import matplotlib.pyplot as plt
import matplotlib.cm as cm
from graphenemodeling.graphene import monolayer as mlg
from scipy.constants import elementary_charge, hbar
eV = elementary_charge
FermiLevel = 0.4 * eV
gamma = 0.012 * eV / hbar
kF = mlg.FermiWavenumber(FermiLevel,model='LowEnergy')
q = np.linspace(1e-2,3,num=200) * kF
w = np.linspace(1e-2,3,num=200) * FermiLevel / hbar
fresnelTM = mlg.FresnelReflection(q,w[:,np.newaxis],gamma,FermiLevel,T=0,
                                    eps1=1,eps2=1,
                                    polarization='TM')
fig, ax = plt.subplots(figsize=(6,6))
ax.imshow(-np.imag(fresnelTM),
            extent=(q[0]/kF,q[-1]/kF,hbar*w[0]/FermiLevel,hbar*w[-1]/FermiLevel),
            origin='lower',aspect='auto',cmap=cm.hot,vmin=-16,vmax=0)
ax.set_xlabel('$q/k_F$')
ax.set_ylabel('$\hbar\omega/E_F$')
ax.set_ylim(0,3)
ax.set_xlim(0,3)
fig.suptitle('Fresnel Reflection Coefficient (TM)')
plt.show()