import kwant

#for plotting
from matplotlib import pyplot

def make_system(a=1, t = 1.0, W = 500, L = 1000):
    lat = kwant.lattice.square(a) #square lattice
    sys = kwant.Builder() #tight binding system
    
    sys[(lat(x,y) for x in range(L) for y in range(W)) ] = 4*t
    sys[lat.neighbors()] = -t
    
    lead = kwant.Builder(kwant.TranslationalSymmetry((-a,0)))
    lead[(lat(0,j) for j in xrange(W))] = 4*t
    lead[lat.neighbors()] = -t
    
    sys.attach_lead(lead)
    sys.attach_lead(lead.reversed())
    return sys
    
#Electron Density
def density(sys,energy,lead_nr):
        wf = kwant.solvers.mumps.wave_function(sys,energy, args = [-.4])
        return (abs(wf(lead_nr))**2).sum(axis=0)
#End Electron Density
def plot_conductance(sys,energies):
    #compute conductance
    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()
    
def main():
    sys = make_system()
    
    kwant.plot(sys,num_lead_cells = 20)
    
    sys = sys.finalized()
    
    #plot_conductance(sys,energies = [0.01 * i for i in xrange(100)])
    #local_dos = kwant.ldos(sys,energy = .3)
    #kwant.plotter.map(sys,local_dos,num_lead_cells = 20, cmap = 'bwr')
    
    #d = density(sys,.2,0)
    #kwant.plotter.map(sys,d,num_lead_cells = 30,cmap = 'bwr')
    
if __name__ == '__main__':
    main()