#!/usr/bin/env python
# coding: utf-8

# In[12]:


get_ipython().run_line_magic('run', '../code/init_mooc_nb.ipy')

import scipy
from matplotlib import cm
from matplotlib import gridspec  
import tinyarray as ta
import scipy.sparse.linalg as sl


sigma0 = np.array([[1, 0], [0, 1]])
sigmax = np.array([[0, 1], [1, 0]])
sigmay = np.array([[0, -1j], [1j, 0]])
sigmaz = np.array([[1, 0], [0, -1]])


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

        
# Onsite and hoppings matrices used for building BHZ model
def onsite(site, par):
    A1, A2, B1, B2, D1, D2, C, M = par.A1, par.A2, par.B1, par.B2, par.D1, par.D2, par.C, par.M
    return (C + 2 * D1 + 4 * D2) * np.kron(sigma0, sigma0) + (M + 2 * B1 + 4 * B2) * np.kron(sigma0, sigmaz)


def hopx(site1, site2, par):
    A1, A2, B1, B2, D1, D2, C, M = par.A1, par.A2, par.B1, par.B2, par.D1, par.D2, par.C, par.M
    return - D2 * np.kron(sigma0, sigma0) - B2 * np.kron(sigma0, sigmaz) + A2 * 0.5j * np.kron(sigmax, sigmax)


def hopy(site1, site2, par):
    A1, A2, B1, B2, D1, D2, C, M = par.A1, par.A2, par.B1, par.B2, par.D1, par.D2, par.C, par.M
    return - D2 * np.kron(sigma0, sigma0) - B2 * np.kron(sigma0, sigmaz) + A2 * 0.5j * np.kron(sigmay, sigmax)


#def hopz(site1, site2, par):
 #   A1, A2, B1, B2, D1, D2, C, M = par.A1, par.A2, par.B1, par.B2, par.D1, par.D2, par.C, par.M
#  return - D1 * np.kron(sigma0, sigma0) - B1 * np.kron(sigma0, sigmaz) + A1 * 0.5j * np.kron(sigmaz, sigmax)

    
def bhz_slab(L, W):
    #lat = kwant.lattice.general(np.identity(3))
    lat = kwant.lattice.square()
    sys = kwant.Builder()
    
    sym0 = kwant.TranslationalSymmetry((1, 0))
    lead0 = kwant.Builder(sym0)
    
    sym1 = kwant.TranslationalSymmetry((-1, 0))
    lead1 = kwant.Builder(sym1)
    
    def shape_lead(pos):
        (x, y) = pos
        return (0 <= y < W)
    
    def shape(pos):
        (x, y) = pos
        return (0 <= x < L) and (0 <= y < W)
    
    def hopping_x(site1, site2, par):
        xt, yt = site1.pos
        xs, ys = site2.pos
        return hopx(site1,site2,par) #* np.exp(-0.5j * par.Bz * (xt - xs) * (yt + ys))
    
    # scattering system
    sys[lat.shape(shape, (0,0))] = lambda site, par: onsite(site, par.scat) - par.mu * np.eye(4)
    sys[kwant.HoppingKind((1,0), lat)] = lambda site1, site2, par: hopx(site1, site2, par.scat)
    sys[kwant.HoppingKind((0,1), lat)] = lambda site1, site2, par: hopy(site1, site2, par.scat)
    #sys[kwant.HoppingKind((0,0,1), lat)] = lambda site1, site2, par: hopz(site1, site2, par.scat)

    # leads
    lead0[lat.shape(shape_lead, (0,0))] = lambda site, par: onsite(site, par.lead) - par.mu_lead * np.eye(4)
    lead0[kwant.HoppingKind((1,0), lat)] = lambda site1, site2, par: hopx(site1, site2, par.lead)
    lead0[kwant.HoppingKind((0,1), lat)] = lambda site1, site2, par: hopy(site1, site2, par.lead)
    #lead0[kwant.HoppingKind((0,0,1), lat)] = lambda site1, site2, par: hopz(site1, site2, par.lead)
    
    lead1[lat.shape(shape_lead, (0,0))] = lambda site, par: onsite(site, par.lead) - par.mu_lead * np.eye(4)
    lead1[kwant.HoppingKind((1,0), lat)] = lambda site1, site2, par: hopx(site1, site2, par.lead)
    lead1[kwant.HoppingKind((0,1), lat)] = lambda site1, site2, par: hopy(site1, site2, par.lead)
    #lead1[kwant.HoppingKind((0,0,1), lat)] = lambda site1, site2, par: hopz(site1, site2, par.lead)
    
    sys.attach_lead(lead0)
    sys.attach_lead(lead1)
    return sys#.finalized() 


# In[13]:


L, W = 10, 10
sys = bhz_slab(L, W)
kwant.plot(sys)


# In[15]:


def site_indexer(wavefunction, sites, n_orbs=4):
    site_idx_map = {site: i for i, site in enumerate(sites)}
    def indexer(site):
        orbital = site_idx_map[site] * n_orbs
        return wavefunction[orbital: orbital + n_orbs]

    return indexer

mu_lead = 0.7
mu_max = .7
L, W = 4, 3
lat = kwant.lattice.square()
sys = bhz_slab(L, W)
fsys = sys.finalized()
par_lead = SimpleNamespace(A1=1.0, A2=1.0, B1=1.0, B2=1.0, C=0.0, D1=0., D2=0., M=-1.0, Bz=Bz)
par = SimpleNamespace(lead=par_lead, scat=par_lead, mu=0.0, mu_lead=0.0)

wf = kwant.solvers.default.wave_function(fsys, energy=1.7, args=[par,])
wavefunction = wf(0)

print wavefunction

#better_wf = site_indexer(wf, fsys.sites)
better_wf = site_indexer(wavefunction, fsys.sites)

# we get an array of length 4, for the 4 orbitals on the site
print better_wf(lat(0, 1))


# In[ ]: