import matplotlib.pyplot as plt
import kwant
import kwant.continuum



syst = kwant.Builder()
print(syst)

a = 10 # angstroms
norbs = 2 # 2 atomic orbitals * 2 spins * particle-hole
at = kwant.lattice.cubic(a, norbs=norbs)
#chain (1D), cubic (3D), triangular, honeycomb, kagome
#print(lat)

ham = ("alpha * (k_x * sigma_x - k_y * sigma_y)"
            "+ (m + beta * kk) * sigma_z"
           "+ (gamma * kk + U) * sigma_0")
subs = {"kk": "k_x**2 + k_y**2"}


ham_discretized = kwant.continuum.discretize(ham, locals=subs, grid=a)
print(ham_discretized)

# dimensions in angstroms
L = 300
W = 50
H = 100

def syst_shape(site):
    (x, y) = site.pos
    return (0 <= x <= L and 0 <= y <= W)

syst.fill(ham_discretized, syst_shape, (0, 0));

lead_left = kwant.Builder(kwant.TranslationalSymmetry((-a, 0)))
lead_right = kwant.Builder(kwant.TranslationalSymmetry((a, 0)))                   

L_lead = L
W_lead = W

def lead_shape(site):
    (x, y) = site.pos
    return (0 <= x <= L_lead and 0 <= y <= W_lead)

lead_left.fill(ham_discretized, lead_shape, (0, 0));
lead_right.fill(ham_discretized, lead_shape, (L, 0));

syst_leads = syst

syst_leads.attach_lead(lead_left)
syst_leads.attach_lead(lead_right)

kwant.plot(syst_leads);

syst = syst.finalized()
syst_leads = syst_leads.finalized()

params = dict(alpha=0.365, beta=0.686, gamma=0.512, m=-0.01, U=-0)


psi = kwant.wave_function(syst, energy=params['m'], params=params)(0)
J = kwant.operator.Current(syst).bind(params=params)
current = sum(J(p) for p in psi)
kwant.plotter.current(syst, current)



kwant.plotter.bands(syst_leads.leads[0], params=params);

energies = []
data = []
for ie in range(100):
    energy = ie * 0.01

    # compute the scattering matrix at a given energy
    smatrix = kwant.smatrix(syst, energy)

    # compute the transmission probability from lead 0 to
    # lead 1
    energies.append(energy)
    data.append(smatrix.transmission(0,1)(0,0))


plt.figure()
plt.plot(energies, data)
plt.xlabel("energy [t]")
plt.ylabel("conductance [e^2/h]")
plt.show()