%matplotlib inline
import numpy as np
from matplotlib import pyplot as plt
from cmath import exp
import math
import kwant
from kwant.digest import gauss
from matplotlib.colors import LogNorm

plt.rcParams['savefig.dpi']=150 #(6.0,4.0)

def hopping(sitei, sitej, phi, salt):
    xi, yi = sitei.pos
    xj, yj = sitej.pos
    return -exp(-0.5j * phi * (xi - xj) * (yi + yj))

def onsite(site, phi, salt):
    return 0.05 * gauss(repr(site), salt) + 4

def make_system(L=60):
    def central_region(pos):
        x, y = pos
        return -2*L < x < 2*L and \
            abs(y) < L - 50 * math.exp(-x**2 / 20**2)

    lat = kwant.lattice.square()
    sys = kwant.Builder()

    sys[lat.shape(central_region, (0, 0))] = onsite
    sys[lat.neighbors()] = hopping

    sym = kwant.TranslationalSymmetry((-1, 0))
    lead = kwant.Builder(sym)
    lead[(lat(0, y) for y in range(-L + 1, L))] = 4
    lead[lat.neighbors()] = hopping

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

    

lat = kwant.lattice.square()
sys = make_system()
fsys = sys.finalized()
lattice_coordinates = kwant.plotter.sys_leads_sites(sys, num_lead_cells=0)
lattice_coordinates = kwant.plotter.sys_leads_pos(sys, lattice_coordinates[0])

#kwant.plot(sys, site_lw=0.05, num_lead_cells=20)

energy = 0.3
phi = 0.045


def plot_bandstructure(flead, momenta):
    bands = kwant.physics.Bands(flead,[phi, ""])
    energies = [bands(k) for k in momenta]

    plt.figure()
    plt.plot(momenta, energies,'b')
    plt.xlabel("momentum [(lattice constant)^-1]")
    plt.ylabel("energy [t]")
    plt.ylim(0,1)
#     plt.xlim(-1,1)
    plt.title('Bottom Contact Band Structure')
    plt.show()
    
momenta = [-np.pi + 0.02/6 * np.pi * i for i in xrange(101*6)]
plot_bandstructure(fsys.leads[0], momenta)

def density(sys, energy, args, lead_nr):
    wf = kwant.wave_function(sys, energy, args)
    return (abs(wf(lead_nr))**2).sum(axis=0)

d = density(fsys, energy, [phi, ""], 0)

#Plotting
img , min ,max = kwant.plotter.mask_interpolate(lattice_coordinates,d,a = 1,oversampling = 3,method = "linear")
border = 0.5 * (max - min) / (np.asarray(img.shape) - 1)
min -= border
max += border
plt.figure(dpi=400)
plt.imshow(img.T, norm = LogNorm(vmin = (1e-4)*np.max(d)), extent=(min[0], max[0], min[1], max[1]), origin='lower', cmap="nipy_spectral") #nipy_spectral
plt.xlabel("X")
plt.ylabel("Y")
plt.title(r"E = "+str(energy)+" , "+r"$ \phi = $"+" "+str(phi))
plt.colorbar()
plt.show()
#end plotting

wf = kwant.wave_function(fsys, energy, [phi, ""])

# Current function based on suggestion from Joseph Weston https://www.mail-archive.com/kwant-discuss%40kwant-project.org/msg00075.html
wf = wf(0) #Get the wave function from the bottom lead
def current(site_i,site_j,E):
    H_ij = fsys.hamiltonian(site_i,site_j,phi,"") # Hamiltonian element between site i and site j
    return -2 * wf[0,site_i].conjugate() * H_ij * wf[0,site_j].imag

sites, lead_sites_slcs = kwant.plotter.sys_leads_sites(fsys, 0) 
n_sys_sites = sum(i[1] is None for i in sites)
                  
sites_pos = kwant.plotter.sys_leads_pos(fsys, sites)

current_density =[] # Initialize Array
for x in range(0, n_sys_sites):
    all_the_neighbors = fsys.graph.out_neighbors(x)
    I_site_i = 0
    for item in all_the_neighbors:
        I_site_i += current(x,item,energy)
    current_density = np.append(current_density, I_site_i)

CD = np.abs(current_density) # Take absolute valute of the current density for plotting

#plotting
img , min ,max = kwant.plotter.mask_interpolate(lattice_coordinates,CD,a = 1,oversampling = 3,method = "linear")
border = 0.5 * (max - min) / (np.asarray(img.shape) - 1)
min -= border
max += border
plt.figure(dpi=400)
plt.imshow(img.T, norm = LogNorm((1e-4)*np.max(CD)), extent=(min[0], max[0], min[1], max[1]), 
                                       origin='lower', cmap="nipy_spectral")
plt.xlabel("X")
plt.ylabel("Y")
plt.title(r"E = "+str(energy)+" , "+r"$ \phi = $"+" "+str(phi))
plt.colorbar()
plt.show()
#end plotting