"""
2DHG TMF transmission and LDOS calculation
DOI: 10.5281/zenodo.14933263
Author: Yik Kheng Lee
"""

import kwant
import numpy as np
from matplotlib import pyplot
import tinyarray
from functools import reduce
from scipy import sparse
from scipy import integrate

# Spin 3/2 matrices
rt3 = np.sqrt(3)
rt2 = np.sqrt(2)
sigma_0 = tinyarray.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
sigma_x = tinyarray.array([[0, 0, rt3, 0], [0, 0, 0, rt3], [rt3, 0, 0, 2], [0, rt3, 2, 0]])/2
sigma_y = tinyarray.array([[0, 0, -rt3, 0], [0, 0, 0, rt3], [rt3, 0, 0, -2], [0, -rt3, 2, 0]])*1j/2
sigma_z = tinyarray.array([[3, 0, 0, 0], [0, -3, 0, 0], [0, 0, 1, 0], [0, 0, 0, -1]])/2

# constants
h_bar = 1.055 * 10**(-34)   # Planck's constant (J.s)
m_e = -9.109 * 10**(-31)    # Hole rest mass (kg)
q_e = 1.602 * 10**(-19)     # Electron charge (C)
mu_B = 9.274 * 10**(-24)    # Bohr magneton (J/T)
j2eV = 6.242 * 10**(18)     # J to eV conversion factor
mu = h_bar**2/m_e

# a: unit cell length 
# W: scattering area width (y)
# L: length between leads (x)
# W1: lead 1 width
# W2: lead 2 width
# Bz: out of plane magnetic field
# lead_l: length in y of QPC gate potential area
# depth: depth of 2DEG from surface gate
# Vg_1, Vg_2: Gate voltage 1 and 2
# U0: Disorder strength
# U_seed: Disorder random seed
# g1, g2, g3: Luttinger parameters
def make_system(a=1, W=5, L=5, L1=5, L2=5, Bz=0, lead_l=0, U0=0, U_seed=0, g1=0, g2=0, g3=0):
    # create square lattice with lattice constant a
    lat = kwant.lattice.square(norbs=4)
    # initiate builder
    syst = kwant.Builder()
    # number of buffer sites in x before lead 1 and after lead 2
    wd = int(4*L1)
    
    # Peierls substitution
    def gauge(site1, site2):
        A1 = (a*1e-9)**2 * ((Bz/2)*(site1.pos[0] - site2.pos[0])*(site1.pos[1] + site2.pos[1]) + trans(site2)*window_f(site2) - trans(site1)*window_f(site1))
        return A1 * (q_e / h_bar)
    
    def gaugex(site1, site2):
        A1 = (a*1e-9)**2 * (Bz/2)*(site1.pos[0] - site2.pos[0])*(site1.pos[1] + site2.pos[1])
        return A1 * (q_e / h_bar)
    
    def gaugey(site1, site2):
        A1 = (a*1e-9)**2 * (Bz/2)*(site1.pos[0] + site2.pos[0])*(site2.pos[1] - site1.pos[1])
        return A1 * (q_e / h_bar)
    
    def trans(site):
        F1 = Bz*site.pos[0]*site.pos[1] 
        return F1
    
    # window when gauge changes
    def window_f(site):
        if site.pos[0] > wd+L1+L+L2+wd-3 or site.pos[0] < 3:
            return 0
        else:
            return 1   

    # Disorder
    np.random.seed(U_seed)
    def disorder():
        U = U0 * (np.random.uniform() - 0.5) * zero_mat
        return U
    
    # Calculate wave functions in z
    # W: width of triangular confinement potential in z (nm)
    # a: unit cell length (nm) (<0.1 for convergence, 0.0005 used for calculations in paper)
    # V0: height of confinement potential 
    def z_confinement(W=100, a=0.1, V0=211):  
        # effective mass factors
        m_f_h = -1/((g1 - 2*g2))
        m_f_l = -1/((g1 + 2*g2))
        
        M = round(W/a - 1)    
        N1 = M    
        N2 = M
        
        L1 = (N1+1)*a
        L2 = (N2+1)*a
        
        z1 = np.empty(0)
        z2 = np.empty(0)
        if N1 > 0:
            z1 = np.linspace(-L1, 0, N1)
        if N2 > 0:
            z2 = np.linspace(a*M, a*M+L2, N2)
        z_middle = np.linspace(0, a*M, M)
        
        z_all = np.concatenate((z1, z_middle, z2))
        
        #Construct the potential well
        if N1 > 0:
            V1 = V0*np.ones(N1)
        else:
            V1 = np.empty(0)
        if N2 > 0:
            V2 = V0*np.ones(N2)
        else:
            V2 = np.empty(0)
        V_middle = np.linspace(0, (2.11*a)*M, M)
        V1D = np.concatenate((V1,V_middle,V2))
        
        N = len(V1D)
        V = sparse.diags(V1D, 0)
        
        # Start finite difference approach
        
        T = sparse.diags([-1, +2, -1], [-1, 0, 1], shape=(N, N))
        H_fact_h = (h_bar**2/(2*m_f_h*m_e*a**2))*(1e3*1e18/q_e)
        H_fact_l = (h_bar**2/(2*m_f_l*m_e*a**2))*(1e3*1e18/q_e)
        
        H_h = (H_fact_h * T + V)
        H_l = (H_fact_l * T + V)
        
        e_val_h, e_vec_h = sparse.linalg.eigsh(H_h, k=1, which="SM")
        e_val_l, e_vec_l = sparse.linalg.eigsh(H_l, k=1, which="SM")
        
        for n in range(np.shape(e_val_h)[0]):
            e_vec_h[:,n] = np.sign(e_vec_h[:,n][N1+1])*e_vec_h[:,n]
            e_vec_l[:,n] = np.sign(e_vec_l[:,n][N1+1])*e_vec_l[:,n]
        
        fH2 = integrate.simpson(e_vec_h[:,0]*e_vec_h[:,0], x=z_all, dx=a)
        fL2 = integrate.simpson(e_vec_l[:,0]*e_vec_l[:,0], x=z_all, dx=a)
        
        def dx(fx,a):
            d1 = np.zeros(np.shape(fx)[0])
            d2 = np.zeros(np.shape(fx)[0])
            
            d1[0] = ((fx[1] - fx[0]))/a
            d2[0] = (fx[0] - 2*fx[1] + fx[2])/a**2
            d1[-1] = ((fx[-1] - fx[-2]))/a
            d2[-1] = (fx[-1 ] - 2*fx[-2] + fx[-3])/a**2 
            for n in range(1,np.shape(fx)[0]-1):
                d1[n] = ((fx[n+1] - fx[n-1])/2)/a
                d2[n] = (fx[n-1] - 2*fx[n] + fx[n+1])/a**2
            
            return d1, d2
        
        dx_h, dx2_h = dx(e_vec_h[:,0], a)
        dx_l, dx2_l = dx(e_vec_l[:,0], a)
        
        f_HL = integrate.simpson(np.conjugate(e_vec_h[:,0])*e_vec_l[:,0], x=z_all, dx=a)/np.sqrt(fH2 * fL2)
        f_LH = integrate.simpson(np.conjugate(e_vec_l[:,0])*e_vec_h[:,0], x=z_all, dx=a)/np.sqrt(fL2 * fH2)
        f_HkL = -1j*integrate.simpson(np.conjugate(e_vec_h[:,0])*dx_l, x=z_all, dx=a)/np.sqrt(fH2 * fL2)
        f_LkH = -1j*integrate.simpson(np.conjugate(e_vec_l[:,0])*dx_h, x=z_all, dx=a)/np.sqrt(fL2 * fH2)
        f_Hk2H = -1*integrate.simpson(np.conjugate(e_vec_h[:,0])*dx2_h, x=z_all, dx=a)/np.sqrt(fH2 * fH2)
        f_Lk2L = -1*integrate.simpson(np.conjugate(e_vec_l[:,0])*dx2_l, x=z_all, dx=a)/np.sqrt(fL2 * fL2)
        
        return f_HL, f_LH, f_HkL, f_LkH, f_Hk2H, f_Lk2L

    f_HL, f_LH, f_HkL, f_LkH, f_Hk2H, f_Lk2L = z_confinement(100, 0.1, 211)
    
    zero_mat = tinyarray.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    P_mat = tinyarray.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    Q_mat = tinyarray.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, -1, 0], [0, 0, 0, -1]])
    M_mat = tinyarray.array([[0, 0, 0, f_HL], [0, 0, 0, 0], [0, f_LH, 0, 0], [0, 0, 0, 0]])
    M_star_mat = tinyarray.array([[0, 0, 0, 0], [0, 0, f_HL, 0], [0, 0, 0, 0], [f_LH, 0, 0, 0]])
    L_mat = tinyarray.array([[0, 0, f_HkL, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, -f_LkH, 0, 0]])
    L_star_mat = tinyarray.array([[0, 0, 0, 0], [0, 0, 0, -f_HkL], [f_LkH, 0, 0, 0], [0, 0, 0, 0]])
    
    mu1 = mu*1e3*1e18/q_e
    # define onsite term
    eps = ((2*mu1/a**2)*(g1*P_mat + g2*Q_mat)
            + (mu1/4)*(g1 - 2*g2)*f_Hk2H*(P_mat+Q_mat)
            + (mu1/4)*(g1 + 2*g2)*f_Lk2L*(P_mat-Q_mat)
            ) 
    
    # define hopping terms
    # nearest x 
    hopx1 = (- (mu1/(2*a**2))*(g1*P_mat + g2*Q_mat)
             + 1j*(rt3/(2*a))*(mu1*g3)*L_mat
             + 1j*(rt3/(2*a))*(mu1*g3)*L_star_mat
             + (rt3/(2*a**2))*(mu1*g2)*M_mat
             + (rt3/(2*a**2))*(mu1*g2)*M_star_mat
             )
    
    # nearest y
    hopy1 = (- (mu1/(2*a**2))*(g1*P_mat + g2*Q_mat)
             + (rt3/(2*a))*(mu1*g3)*L_mat
             - (rt3/(2*a))*(mu1*g3)*L_star_mat
             - (rt3/(2*a**2))*(mu1*g2)*M_mat
             - (rt3/(2*a**2))*(mu1*g2)*M_star_mat
             )
        
    # diagonal upwards
    hopd_u = (- rt3/(4*a**2)*1j*mu1*g3*M_mat 
              + rt3/(4*a**2)*1j*mu1*g3*M_star_mat
              )
        
    # diagonal downwards
    hopd_d = (rt3/(4*a**2)*1j*mu1*g3*M_mat 
              - rt3/(4*a**2)*1j*mu1*g3*M_star_mat
              )
    
    # set onsite energies for all sites in system
    for i in range(wd+L1+L+L2+wd):
        for j in range(lead_l, lead_l + W):
            syst[lat(i,j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    # Square leads
    for i in range(wd, wd+L1):
        for j in range(lead_l):
            syst[lat(i,j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    for i in range(wd+L1+L, wd+L1+L+L2):
        for j in range(lead_l+1):
            syst[lat(i,j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    # Trumpet leads
    for j in range(1, lead_l):
        for i in range(wd-j, wd):
            syst[lat(i,j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
        for i in range(wd+L1, wd+L1+j):
            syst[lat(i,j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
        for i in range(wd+L1+L-j, wd+L1+L):
            syst[lat(i,j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
        for i in range(wd+L1+L+L2, wd+L1+L+L2+j):
            syst[lat(i,j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    
    
    # set hopping energies, kwant auto applies forward and backward hoppings
    for i in range(wd+L1+L+L2+wd):
        for j in range(lead_l, lead_l + W):
            # hopping in x
            if i > 0:
                syst[lat(i - 1, j), lat(i, j)] = hopx1 * np.exp(1j * gauge(lat(i - 1, j), lat(i, j)))
            if j > lead_l:
                syst[lat(i, j - 1), lat(i, j)] = hopy1 * np.exp(1j * gauge(lat(i, j - 1), lat(i, j)))
            # diagonal hopping
            if i > 0 and j > lead_l:
                syst[lat(i - 1, j - 1), lat(i, j)] = hopd_u * np.exp(1j * gauge(lat(i - 1, j - 1), lat(i, j)))
                syst[lat(i - 1, j), lat(i, j - 1)] = hopd_d * np.exp(1j * gauge(lat(i - 1, j), lat(i, j - 1)))
    # Square leads
    for i in range(wd, wd+L1):
        for j in range(lead_l + 1):
            if i > wd:
                syst[lat(i - 1, j), lat(i, j)] = hopx1 * np.exp(1j * gauge(lat(i - 1, j), lat(i, j)))
            if j > 0:
                syst[lat(i, j - 1), lat(i, j)] = hopy1 * np.exp(1j * gauge(lat(i, j - 1), lat(i, j)))
            if i > wd and j > 0:
                syst[lat(i - 1, j - 1), lat(i, j)] = hopd_u * np.exp(1j * gauge(lat(i - 1, j - 1), lat(i, j)))
                syst[lat(i - 1, j), lat(i, j - 1)] = hopd_d * np.exp(1j * gauge(lat(i - 1, j), lat(i, j - 1)))
    for i in range(wd+L1+L, wd+L1+L+L2):
        for j in range(lead_l + 1):
            if i > wd+L1+L:
                syst[lat(i - 1, j), lat(i, j)] = hopx1 * np.exp(1j * gauge(lat(i - 1, j), lat(i, j)))
            if j > 0:
                syst[lat(i, j - 1), lat(i, j)] = hopy1 * np.exp(1j * gauge(lat(i, j - 1), lat(i, j)))
            if i > wd+L1+L and j > 0:
                syst[lat(i - 1, j - 1), lat(i, j)] = hopd_u * np.exp(1j * gauge(lat(i - 1, j - 1), lat(i, j)))
                syst[lat(i - 1, j), lat(i, j - 1)] = hopd_d * np.exp(1j * gauge(lat(i - 1, j), lat(i, j - 1)))
    # Trumpet leads
    for j in range(1, lead_l + 1):
        for i in range(wd - j, wd + 1):
            if i > wd - j:
                syst[lat(i - 1, j), lat(i, j)] = hopx1 * np.exp(1j * gauge(lat(i - 1, j), lat(i, j)))
                if j > 1:
                    syst[lat(i, j - 1), lat(i, j)] = hopy1 * np.exp(1j * gauge(lat(i, j - 1), lat(i, j)))                   
                    # diagonal hopping
                syst[lat(i - 1, j), lat(i, j - 1)] = hopd_d * np.exp(1j * gauge(lat(i - 1, j), lat(i, j - 1)))
            if i > wd - j + 1 and j > 1:
                syst[lat(i - 1, j - 1), lat(i, j)] = hopd_u * np.exp(1j * gauge(lat(i - 1, j - 1), lat(i, j)))                    
        for i in range(wd+L1, wd+L1+j):
            syst[lat(i - 1, j), lat(i, j)] = hopx1 * np.exp(1j * gauge(lat(i - 1, j), lat(i, j)))
            if j > 1:
                if i < wd+L1+j-1:
                    syst[lat(i, j - 1), lat(i, j)] = hopy1 * np.exp(1j * gauge(lat(i, j - 1), lat(i, j)))
                # diagonal hopping
                    syst[lat(i - 1, j), lat(i, j - 1)] = hopd_d * np.exp(1j * gauge(lat(i - 1, j), lat(i, j - 1)))
            if i < wd+L1+j:
                syst[lat(i - 1, j - 1), lat(i, j)] = hopd_u * np.exp(1j * gauge(lat(i - 1, j - 1), lat(i, j)))                 
        for i in range(wd+L1+L-j, wd+L1+L+1):
            if i > wd+L1+L-j:
                syst[lat(i - 1, j), lat(i, j)] = hopx1 * np.exp(1j * gauge(lat(i - 1, j), lat(i, j)))
                if j > 1:
                    syst[lat(i, j - 1), lat(i, j)] = hopy1 * np.exp(1j * gauge(lat(i, j - 1), lat(i, j)))                   
                    # diagonal hopping
                syst[lat(i - 1, j), lat(i, j - 1)] = hopd_d * np.exp(1j * gauge(lat(i - 1, j), lat(i, j - 1)))
            if i > wd+L1+L-j + 1 and j > 1:
                syst[lat(i - 1, j - 1), lat(i, j)] = hopd_u * np.exp(1j * gauge(lat(i - 1, j - 1), lat(i, j)))                    
        for i in range(wd+L1+L+L2, wd+L1+L+L2+j):
            syst[lat(i - 1, j), lat(i, j)] = hopx1 * np.exp(1j * gauge(lat(i - 1, j), lat(i, j)))
            if j > 1:
                if i < wd+L1+L+L2+j-1:
                    syst[lat(i, j - 1), lat(i, j)] = hopy1 * np.exp(1j * gauge(lat(i, j - 1), lat(i, j)))
                # diagonal hopping
                    syst[lat(i - 1, j), lat(i, j - 1)] = hopd_d * np.exp(1j * gauge(lat(i - 1, j), lat(i, j - 1)))
            if i < wd+L1+L+L2+j:
                syst[lat(i - 1, j - 1), lat(i, j)] = hopd_u * np.exp(1j * gauge(lat(i - 1, j - 1), lat(i, j))) 
    
    ### create leads

    d_lead1 = kwant.Builder(kwant.TranslationalSymmetry((0, -1)))
    for i in range(wd, wd+L1):
        # onsite energies
        d_lead1[lat(i, 0)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
        d_lead1[lat(i, 1)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    for i in range(wd, wd+L1):
        # hopping in x
        if i > wd:
            d_lead1[lat(i - 1, 0), lat(i, 0)] = hopx1 * np.exp(1j * gauge(lat(i - 1, 0), lat(i, 0)))
        # hopping in y
        d_lead1[lat(i, 0), lat(i, 1)] = hopy1 * np.exp(1j * gauge(lat(i, 0), lat(i, 1)))
        # diagonal hopping
        if i > wd:
            d_lead1[lat(i - 1, 0), lat(i, 1)] = hopd_u * np.exp(1j * gauge(lat(i - 1, 0), lat(i, 1)))
            d_lead1[lat(i - 1, 1), lat(i, 0)] = hopd_d * np.exp(1j * gauge(lat(i - 1, 1), lat(i, 0)))
    # attach lead to channel
    syst.attach_lead(d_lead1)
    
    d_lead2 = kwant.Builder(kwant.TranslationalSymmetry((0, -1)))
    for i in range(wd+L1+L, wd+L1+L+L2):
        # onsite energies
        d_lead2[lat(i, 0)] =  eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
        d_lead2[lat(i, 1)] =  eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    for i in range(wd+L1+L, wd+L1+L+L2):
        # hopping in x
        if i > wd+L1+L:
            d_lead2[lat(i - 1, 0), lat(i, 0)] = hopx1 * np.exp(1j * gauge(lat(i - 1, 0), lat(i, 0)))
        # hopping in y
        d_lead2[lat(i, 0), lat(i, 1)] = hopy1 * np.exp(1j * gauge(lat(i, 0), lat(i, 1)))
        # diagonal hopping
        if i > wd+L1+L:
            d_lead2[lat(i - 1, 0), lat(i, 1)] = hopd_u * np.exp(1j * gauge(lat(i - 1, 0), lat(i, 1)))
            d_lead2[lat(i - 1, 1), lat(i, 0)] = hopd_d * np.exp(1j * gauge(lat(i - 1, 1), lat(i, 0)))
    # attach lead to channel
    syst.attach_lead(d_lead2)

    u_lead = kwant.Builder(kwant.TranslationalSymmetry((0, 1)))
    for i in range(1,wd+L1+L+L2+wd-1):
        # onsite energies
        u_lead[lat(i, 0)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
        u_lead[lat(i, 1)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    for i in range(1,wd+L1+L+L2+wd-1):
        # hopping in x
        if i > 1:
            u_lead[lat(i - 1, 0), lat(i, 0)] = hopx1 * np.exp(1j * gauge(lat(i - 1, lead_l + W - 1), lat(i, lead_l + W - 1)))
        # hopping in y
        u_lead[lat(i, 0), lat(i, 1)] = hopy1 * np.exp(1j * gauge(lat(i, lead_l + W - 1), lat(i, lead_l + W)))
        # diagonal hopping
        if i > 1:
            u_lead[lat(i - 1, 0), lat(i, 1)] = hopd_u * np.exp(1j * gauge(lat(i - 1, lead_l + W - 1), lat(i, lead_l + W)))
            u_lead[lat(i - 1, 1), lat(i, 0)] = hopd_d * np.exp(1j * gauge(lat(i - 1, lead_l + W), lat(i, lead_l + W - 1)))
    # attach lead to channel
    syst.attach_lead(u_lead)

    l_lead = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))
    for j in range(lead_l,lead_l + W):
        # onsite energies
        l_lead[lat(0, j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
        l_lead[lat(-1, j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    for j in range(lead_l,lead_l + W):
        # hopping in y
        if j > lead_l:
            l_lead[lat(0, j - 1), lat(0, j)] = hopy1 * np.exp(1j * gauge(lat(0, j - 1), lat(0, j)))
        # hopping in x
        l_lead[lat(0, j), lat(1, j)] = hopx1 * np.exp(1j * gauge(lat(0, j), lat(1, j)))
        # diagonal hopping
        if j > lead_l:
            l_lead[lat(0, j - 1), lat(1, j)] = hopd_u * np.exp(1j * gauge(lat(0, j - 1), lat(1, j)))
            l_lead[lat(0, j), lat(1, j - 1)] = hopd_d * np.exp(1j * gauge(lat(0, j), lat(1, j - 1)))
    # attach lead to channel
    syst.attach_lead(l_lead)
    
    r_lead = kwant.Builder(kwant.TranslationalSymmetry((1, 0)))
    for j in range(lead_l,lead_l + W):
        # onsite energies
        r_lead[lat(0, j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
        r_lead[lat(-1, j)] = eps - 2*1.2*mu_B*Bz*sigma_z*j2eV*1e3
    for j in range(lead_l,lead_l + W):
        # hopping in y
        if j > lead_l:
            r_lead[lat(0, j - 1), lat(0, j)] = hopy1 * np.exp(1j * gauge(lat(wd+L1+L+L2+wd-1, j - 1), lat(wd+L1+L+L2+wd-1, j)))
        # hopping in x
        r_lead[lat(0, j), lat(1, j)] = hopx1 * np.exp(1j * gauge(lat(wd+L1+L+L2+wd-1, j), lat(wd+L1+L+L2+wd, j)))
        # diagonal hopping
        if j > lead_l:
            r_lead[lat(0, j - 1), lat(1, j)] = hopd_u * np.exp(1j * gauge(lat(wd+L1+L+L2+wd-1, j - 1), lat(wd+L1+L+L2+wd, j)))
            r_lead[lat(0, j), lat(1, j - 1)] = hopd_d * np.exp(1j * gauge(lat(wd+L1+L+L2+wd-1, j), lat(wd+L1+L+L2+wd, j - 1)))
    # attach lead to channel
    syst.attach_lead(r_lead)
    
    return syst 

l_nm = 50   # lead width (nm)
a_nm = 2    # unit cell length in (nm)
a_act = a_nm * 10**(-9) # unit cell length (m)
L1 = int(l_nm/a_nm - 1) # lead 1 width (number of cells)
L2 = L1     # lead 2 width (number of cells)
L_nm = 200 # lead separation (nm) 
L = int(L_nm/a_nm - 1)  # lead separation (x) (number of cells)
lead_l = L1 # lead lengths (number of cells)
W = L + lead_l # Scattering area width (y) (number of cells)   

Bz = -2   # Out of plane B (T)
        
U0 = 0.0    # Disorder strength
U_seed = 0  # Disorder seed

enrg = -6.8   # energy in meV

# Luttinger parameters
g1 = 6.85
g2 = 2.10
g3 = 2.90
    
syst = make_system(a=a_nm, W=W, L=L, L1=L1, L2=L2, Bz=Bz, lead_l=lead_l, U0=U0, U_seed=U_seed, g1=g1, g2=g2, g3=g3)
syst = syst.finalized()
kwant.plot(syst)

#####
# Calculate transmission
#####
lead_out = 1
lead_in = 0
gf_a = kwant.greens_function(sys=syst, energy=enrg, out_leads=[lead_out], in_leads=[lead_in])
gf = gf_a.submatrix(lead_out, lead_in)
factors = []
for lead, gf2 in ((lead_out, gf), (lead_in, gf.conj().T)):
    possible_se = gf_a.lead_info[lead]
    factors.append(1j * (possible_se - possible_se.conj().T))
    factors.append(gf2)

# Create and resize spin projection matrices                          
Hxu = np.array([[-rt2/4, -rt2/4, -rt3/(2*rt2), -rt3/(2*rt2)]])  # spin-x +3/2
Hxd = np.array([[rt2/4, -rt2/4, -rt3/(2*rt2), rt3/(2*rt2)]])    # spin-x -3/2
Lxu = np.array([[rt3/(2*rt2), -rt3/(2*rt2), rt2/4, -rt2/4]])    # spin-x +1/2
Lxd = np.array([[-rt3/(2*rt2), -rt3/(2*rt2), rt2/4, rt2/4]])    # spin-x -1/2  
Hxu_mat = [np.kron((np.eye(int(len(factors[0])/4))),Hxu.T @ Hxu)]
Hxd_mat = [np.kron((np.eye(int(len(factors[0])/4))),Hxd.T @ Hxd)]
Lxu_mat = [np.kron((np.eye(int(len(factors[0])/4))),Lxu.T @ Lxu)]
Lxd_mat = [np.kron((np.eye(int(len(factors[0])/4))),Lxd.T @ Lxd)]

Hyu = np.array([[-rt2/4, (rt2/4)*1j, -rt3/(2*rt2)*1j, rt3/(2*rt2)]])  # spin-y +3/2
Hyd = np.array([[(rt2/4)*1j, -rt2/4, rt3/(2*rt2), -rt3/(2*rt2)*1j]])  # spin-y -3/2
Lyu = np.array([[rt3/(2*rt2), rt3/(2*rt2)*1j, (rt2/4)*1j, (rt2/4)]])  # spin-y +1/2
Lyd = np.array([[rt3/(2*rt2)*1j, rt3/(2*rt2), rt2/4, (rt2/4)*1j]])    # spin-y -1/2
Hyu_mat = [np.kron((np.eye(int(len(factors[0])/4))),Hyu.T @ Hyu)]
Hyd_mat = [np.kron((np.eye(int(len(factors[0])/4))),Hyd.T @ Hyd)]
Lyu_mat = [np.kron((np.eye(int(len(factors[0])/4))),Lyu.T @ Lyu)]
Lyd_mat = [np.kron((np.eye(int(len(factors[0])/4))),Lyd.T @ Lyd)]

Hzu = np.array([[1, 0, 0, 0]])  # spin-z +3/2
Hzd = np.array([[0, 1, 0, 0]])  # spin-z -3/2
Lzu = np.array([[0, 0, 1, 0]])  # spin-z +1/2
Lzd = np.array([[0, 0, 0, 1]])  # spin-z -1/2      
Hzu_mat = [np.kron((np.eye(int(len(factors[0])/4))),Hzu.T @ Hzu)]
Hzd_mat = [np.kron((np.eye(int(len(factors[0])/4))),Hzd.T @ Hzd)]
Lzu_mat = [np.kron((np.eye(int(len(factors[0])/4))),Lzu.T @ Lzu)]
Lzd_mat = [np.kron((np.eye(int(len(factors[0])/4))),Lzd.T @ Lzd)]

# Spin projected transmissions           
tt = np.trace(reduce(np.dot, factors)).real             # Total
tx1 = np.trace(reduce(np.dot, Hxu_mat + factors)).real  # spin-x +3/2
tx2 = np.trace(reduce(np.dot, Hxd_mat + factors)).real  # spin-x -3/2
tx3 = np.trace(reduce(np.dot, Lxu_mat + factors)).real  # spin-x +1/2
tx4 = np.trace(reduce(np.dot, Lxd_mat + factors)).real  # spin-x -1/2
ty1 = np.trace(reduce(np.dot, Hyu_mat + factors)).real  # spin-y +3/2
ty2 = np.trace(reduce(np.dot, Hyd_mat + factors)).real  # spin-y -3/2
ty3 = np.trace(reduce(np.dot, Lyu_mat + factors)).real  # spin-y +1/2
ty4 = np.trace(reduce(np.dot, Lyd_mat + factors)).real  # spin-y -1/2
tz1 = np.trace(reduce(np.dot, Hzu_mat + factors)).real  # spin-z +3/2
tz2 = np.trace(reduce(np.dot, Hzd_mat + factors)).real  # spin-z -3/2
tz3 = np.trace(reduce(np.dot, Lzu_mat + factors)).real  # spin-z +1/2
tz4 = np.trace(reduce(np.dot, Lzd_mat + factors)).real  # spin-z -1/2

#####
# Plot LDOS
#####
wf = kwant.wave_function(syst, energy=enrg)
psi = wf(1) 
rho_0 = kwant.operator.Density(syst)
dens_0 = sum(rho_0(p) for p in psi)
fig  = kwant.plotter.map(syst, dens_0, num_lead_cells=3, method="nearest", oversampling=3, colorbar=True, background='white')
ax = pyplot.gca()
import matplotlib.ticker as ticker
ticks_x = ticker.FuncFormatter(lambda x, pos: '${0:.1f}$'.format(x*a_nm/1000))
ax.xaxis.set_major_formatter(ticks_x)
ticks_y = ticker.FuncFormatter(lambda x, pos: '${0:.1f}$'.format(x*a_nm/1000))
ax.yaxis.set_major_formatter(ticks_y)
ax.tick_params(axis='both', which='major', labelsize=15)
ax.set_aspect('equal')
ax.set_xlabel(r'$x$ [$\mathrm{\mu}$m]', fontsize=15)
ax.set_ylabel(r'$y$ [$\mathrm{\mu}$m]', fontsize=15)
pyplot.tight_layout()