import kwant
import numpy as np

sigma_0 = np.array([[1., 0.], [0., 1.]])
sigma_x = np.array([[0., 1.], [1., 0.]])
sigma_p = np.array([[0.,1.],[0.,0.]])
sigma_m = np.array([[0.,0.],[1.,0.]])
sigma_y = np.array([[0., -1j], [1j, 0.]])
sigma_z = np.array([[1., 0.], [0., -1.]])

def make_system_triv(W=5,L=5):
    def onsite_supra(site, t, mu, delt, phi):
        return (4*t - mu) 
    def onsite_normal(site, t, mu, delt, phi):
        return (4*t - mu)
    def hop_normal_x(site1, site2, t, mu, delt, phi):
        return -t * np.exp(1j*phi)
    def hop_normal_y(site1, site2, t, mu, delt, phi):
        return -t 
    def hop_supra(site1, site2, t, mu, delt, phi):
        return -t
    
    lat = kwant.lattice.square(1)
    sys = kwant.Builder()

    sys[(lat(x, y) for x in range(L) for y in range(W))] = onsite_normal
    sys[kwant.builder.HoppingKind((1, 0), lat, lat)] = hop_normal_x
    sys[kwant.builder.HoppingKind((0, 1), lat, lat)] = hop_normal_y

    lead = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))
    lead[(lat(0, j) for j in xrange(W))] = onsite_supra
    lead[kwant.builder.HoppingKind((1, 0), lat, lat)] = hop_supra
    lead[kwant.builder.HoppingKind((0, 1), lat, lat)] = hop_supra

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

    sys = sys.finalized()
    return sys

def make_system(W=5,L=5):
    def onsite_supra(site, t, mu, delt, phi):
        return (4*t - mu) *sigma_z + delt*sigma_x
    def onsite_normal(site, t, mu, delt, phi):
        return (4*t - mu) *sigma_z
    def hop_normal_x(site1, site2, t, mu, delt, phi):
        return -t * np.exp(1j*phi) *sigma_z
    def hop_normal_y(site1, site2, t, mu, delt, phi):
        return -t *sigma_z
    def hop_supra(site1, site2, t, mu, delt, phi):
        return -t *sigma_z
    
    lat = kwant.lattice.square(1)
    sys = kwant.Builder()

    sys[(lat(x, y) for x in range(L) for y in range(W))] = onsite_normal
    sys[kwant.builder.HoppingKind((1, 0), lat, lat)] = hop_normal_x
    sys[kwant.builder.HoppingKind((0, 1), lat, lat)] = hop_normal_y

    lead = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))
    lead[(lat(0, j) for j in xrange(W))] = onsite_supra
    lead[kwant.builder.HoppingKind((1, 0), lat, lat)] = hop_supra
    lead[kwant.builder.HoppingKind((0, 1), lat, lat)] = hop_supra

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

    sys = sys.finalized()
    return sys

fsys = make_system()
l1,l2 = fsys.leads

parameters = [1.,0.,0.,0.]
pro, sta = l1.modes(0.,parameters);
print np.shape(pro.momenta)
#returns (0,)
wf = kwant.wave_function(fsys, 0., parameters)
print np.shape(wf(0))
#returns (0, 50)

parameters = [1.,1.,0.,0.]
pro, sta = l1.modes(0.,parameters);
print np.shape(pro.momenta)
#returns (4,)
wf = kwant.wave_function(fsys, 0., parameters)
print np.shape(wf(0))
#returns (2, 50)

parameters = [1.,0.,1.,0.]
pro, sta = l1.modes(0.,parameters);
print np.shape(pro.momenta)
#returns (0,)
wf = kwant.wave_function(fsys, 0., parameters)
print np.shape(wf(0))
#returns (0, 50)