import matplotlib.pyplot as plt
import numpy as np
import kwant
import scipy.sparse.linalg as sla

class SimpleNamespace(object):
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

# System creation
def make_system(a, t0, S, plot):  # (lattice constant in nm, hopping parameter in eV, square size in nm)
    lat = kwant.lattice.square(a)  # Create a honeycomb lattice
    sys = kwant.Builder()  # Start the builder


    def channel(pos):  # Create all the points in region bounded by
        x, y = pos  # Lattice coordinates
        return -(5 * S) / 2 <= x <= (5 * S) / 2 and 0 <= y <= (5* S)

    def potential(site, par):
        (x, y) = site.pos
        if x > -par.L/2 and x < +par.L/2:
            return par.U_0
        else:
            return par.E1

    def onsite_channel(site, par):
        return 4*t0 + potential(site, par)

    def leadshape0(pos):  # Left lead
        x, y = pos
        return 0 <= y <= 5 * S

    def leadshape1(pos):  # Right lead
        x, y = pos
        return 0 <= y <= 5 * S

    def onsite_lead0(site, par):  # Left lead
        return 4*t0 + par.E1  # + #par.disorder*kwant.digest.uniform(repr(site),salt = par.salt)

    def onsite_lead1(site, par):  # Right lead
        return 4*t0 +par.E2  # + #par.disorder*kwant.digest.uniform(repr(site),salt = par.salt)

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

    sys[lat.shape(channel, (0, 0))] = onsite_channel  # Generate all the coordinates and assign their onsite energy
    sys[lat.neighbors()] = hopping  # Assign hoppings for each atom generated
    #### Define the leads. ####
    # left lead
    # Translational Symmetry of Leads 0,1,....
    sym0 = kwant.TranslationalSymmetry((-a, 0))
    sym1 = kwant.TranslationalSymmetry((+a, 0))

    # End of Translational Symmetry of Leads 0,1,...

    lead0 = kwant.Builder(sym0)
    lead0[lat.shape(leadshape0, (0.5, 0.5))] = onsite_lead0
    lead0[lat.neighbors()] = hopping

    lead1 = kwant.Builder(sym1)
    lead1[lat.shape(leadshape1, (0.5, 0.5))] = onsite_lead1
    lead1[lat.neighbors()] = hopping

    sys.attach_lead(lead0)  # Left
    sys.attach_lead(lead1)  # Right
    kwant.plot(sys,show=False)
    return sys

a_s = 1.
t_s = 1.
S = 5.

par = SimpleNamespace(phi=0.5,
                      E1=-0.5,
                      E2=-0.5,
                      U_0=-0.7,
                      L=2*S,
                      k_val = 0.5)

sys1 = make_system(a_s, t_s, S, "don't plot")
D = kwant.plotter.sys_leads_sites(sys1, num_lead_cells=0)
D = kwant.plotter.sys_leads_pos(sys1, D[0]) # Lattice coordinates
fsys = sys1.finalized()

plt.show()