import kwant
from matplotlib import pyplot
import tinyarray
from numpy import sqrt
sigma_0 = tinyarray.array([[1, 0], [0, 1]])
Sigma_z = tinyarray.array([[1, 0], [0, -1]])


sys=kwant.Builder()
lat=kwant.lattice.honeycomb()
a,b= lat.sublattices
lx,ly=10,7


def Rectangle(pos,lx=lx,ly=ly):
    x,y=pos
    return 0<=x<lx and 1<=y<ly
           
def Hop_Value(site1,site2,Lambda=1,t2=1):
    return 1j*t2*Lambda/(3.*sqrt(3))*Sigma_z



t1=1
sys[lat.shape(Rectangle,(0,1))]=0*sigma_0   #even with 0 onsite potential, we need to use the zero matrix and not a scalar.
sys[lat.neighbors()]=t1*sigma_0

#Here we chose the clockwise Vij which are +1 (as retuned by Hop_value). The other ones (-1) will be given directly 
#by the hermeticity of the Hamiltonian
hoppings1 = (((1, 0), a, a), ((-1, 1), a, a), ((0, -1), a, a))    
hoppings2 = (((1, 0), b, b), ((-1, 1), b, b), ((0, -1), b, b))
sys[[kwant.builder.HoppingKind(*hopping) for hopping in hoppings1]] = Hop_Value
sys[[kwant.builder.HoppingKind(*hopping) for hopping in hoppings2]] = Hop_Value




def lead_shape(pos,ly=ly):
    x,y=pos
    return    1<=y<=ly
    
sym= kwant.TranslationalSymmetry(lat.a.vec((-1, 0)))
lead = kwant.Builder(sym)
lead[lat.shape(lead_shape, (0, 1))] = 0*sigma_0
lead[lat.neighbors()]=t1*sigma_0
lead[[kwant.builder.HoppingKind(*hopping) for hopping in hoppings1]] = Hop_Value
lead[[kwant.builder.HoppingKind(*hopping) for hopping in hoppings2]] = Hop_Value

sys.attach_lead(lead)
sys.attach_lead(lead.reversed())

def family_colors(site):
        return 0 if (site.family == a)  else 1
def hopping_lw(site1, site2):
        return 0.04 if site1.family == site2.family else 0.1
def hopping_color(site1,site2):
         return 'g' if site1.family==site2.family else 'g'
        

kwant.plot(sys,site_color=family_colors,site_lw=0.1, hop_lw=hopping_lw,hop_color=hopping_color,colorbar=False)

sys=sys.finalized()

def plot_conductance(sys, energies):
    
    data = []
    for energy in energies:
        smatrix = kwant.smatrix(sys, energy)
        data.append(smatrix.transmission(1, 0))

    pyplot.figure()
    pyplot.plot(energies, data)
    pyplot.xlabel("energy [t]")
    pyplot.ylabel("conductance [e^2/h]")
    pyplot.show()
energies=[-3+i*0.02 for i in range(300)]
plot_conductance(sys,energies)