import kwant
import matplotlib.pyplot as pyplot
from scipy import misc
import numpy as np

# shape.png is an image containing the system shape colored in black.  Here we
# read the image from this file, create an empty scattering region, and add to
# it the sites with coordinates of pixels that we encounter in the image.
im = misc.imread('shape.png')[:, :, 0]
sys = kwant.Builder()
lat = kwant.lattice.square()
for pos in np.argwhere(im != 255):
    sys[lat(pos[1], -pos[0])] = 4

def wire_shape(period, center, r, ymin, ymax):
    """Return sites from a bounded segment of an infinite wire."""
    direction = np.array(period / np.linalg.norm(period))
    def wire_shape(pos):
        rel_pos = center - pos
        projection = direction * np.dot(direction, rel_pos) - rel_pos
        return sum(projection**2) <= r**2 and ymin < pos[1] < ymax
    return lat.shape(wire_shape, center)

# Directions of the leads, positions of the centers of the lead segments to be
# attached, and the boundaries of these segments in y-directions.
wire_dir = [(-1., -1.), (3., -4.), (2., 4.)]
wire_center = [(166, -236), (216, -106), (246, -86)]
ylims = [(-400, -200), (-440, -100), (-100, 60)]

# Add segments of leads to the scattering region, and attach leads themselves.
for direction, center, ylim in zip(wire_dir, wire_center, ylims):
    sys[wire_shape(direction, np.array(center), 50, *ylim)] = 4
    lead = kwant.Builder(kwant.TranslationalSymmetry(direction))
    lead[lat.wire(center, 50)] = 4
    lead[lat.neighbors()] = -1
    sys.attach_lead(lead)

# Finally add hoppings to the system, get its wave function, and plot.
sys[lat.neighbors()] = -1
sys = sys.finalized()
wfs = kwant.wave_function(sys, energy=0.14)(0)[0]
kwant.plotter.map(sys, abs(wfs), oversampling=1, cmap="PuBu")