import kwant

import matplotlib.pyplot as plt

sys = kwant.Builder()

a = 1
lat = kwant.lattice.square(a)

t = 1.0
W = 60
L = 30

# Define the scattering region

for i in range(L):
    for j in range(W):
        # On-site Hamiltonian
        sys[lat(i, j)] = 4 * t


        if i > 10 and i < 20:
            sys[lat(i, j)] += 6

        # Hopping in y-direction
        if j > 0:
            sys[lat(i, j), lat(i, j - 1)] = -t

        # Hopping in x-direction
        if i > 0:
            sys[lat(i, j), lat(i - 1, j)] = -t

sym_left_lead = kwant.TranslationalSymmetry((-a, 0))
left_lead = kwant.Builder(sym_left_lead)

for j in range(W):
    left_lead[lat(0, j)] = 4 * t
    if j > 0:
        left_lead[lat(0, j), lat(0, j - 1)] = -t
    left_lead[lat(1, j), lat(0, j)] = -t

sys.attach_lead(left_lead)

sym_right_lead = kwant.TranslationalSymmetry((a, 0))
right_lead = kwant.Builder(sym_right_lead)

for j in range(W):
    right_lead[lat(0, j)] = 4 * t
    if j > 0:
        right_lead[lat(0, j), lat(0, j - 1)] = -t
    right_lead[lat(1, j), lat(0, j)] = -t

sys.attach_lead(right_lead)

sys = sys.finalized()

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

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

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

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

plt.figure()
plt.plot(energies, num_prop, '.')
plt.xlabel("energy [t]")
plt.ylabel("num_prop")
plt.show()