import numpy as np
import matplotlib.pyplot as plt
import scipy as sp
import scipy.constants

import kwant
import sys
sys.path.append('./modules/')
import wraparound

# -------------------------------------------------------------------
# Constants
qe = sp.constants.value("elementary charge") # unit: C
me = sp.constants.value("electron mass")/qe*1e-18 #unit: eV*s^2/nm^2
hP = sp.constants.value("Planck constant in eV s") #unit: eV*s
hbar = hP/(2*sp.pi) #unit: eV*s
# -------------------------------------------------------------------

a = 0.03
V0 = 10.0

W = 20
L = 50

# Hole in the rectuangular potential barrier
W_hole = 5
V_hole = 5.0
hole_y = 9

t = hbar**2/(2*me*a**2) # units: eV

lat = kwant.lattice.square(a)

def onsite_potential(site):
    (x,y) = site.pos
    if np.abs(y-hole_y*a) <= W_hole/2*a:
        return V_hole
    return V0

# Infinite potential plane in y direction
sys = kwant.Builder(kwant.TranslationalSymmetry(lat.vec((0, W))))
sys[(lat(i,j) for i in range(L) for j in range(W))] = lambda p: 4*t+onsite_potential(p)
sys[lat.neighbors(1)] = -t

sys = wraparound.wraparound(sys)

lead = kwant.Builder(kwant.TranslationalSymmetry((-a, 0)))
lead[(lat(0,j) for j in range(W))] = 4*t
lead[lat.neighbors(1)] = -t
sys.attach_lead(lead)
sys.attach_lead(lead.reversed())

kwant.plot(sys)

kwant.plot(sys, site_color=onsite_potential, cmap='jet', site_symbol='s', site_size=0.4, fig_size=(10, 6))

sys = sys.finalized()

# -------------------------------------------------------
# Calculation
ky = 1.0

energies = np.arange(6.0, 12.0, 0.04)
transmission = []
for energy in energies:
    smatrix = kwant.smatrix(sys, energy, [ky])
    transmission.append(smatrix.transmission(1, 0))
# -------------------------------------------------------

# Plot transmission
plt.plot(energies, transmission, '.')
plt.yscale('log')
plt.show()