# -*- coding: utf-8 -*-
"""
Created on Tue Aug 14 14:22:18 2020

@author: martifja
"""

import kwant, tinyarray
from math import pi, sqrt
from matplotlib import pyplot as plt
import numpy as np

#Specifying 2x2 Pauli matrices
s_x = tinyarray.array([[0,1],[1,0]])
s_y = tinyarray.array([[0,1j],[1j,0]])
s_z = tinyarray.array([[1,0],[0,-1]])
s_0 = tinyarray.array([[1,0],[0,1]])

#Fix hopping, Fermi energi, and superconducting gap
t = 1
E_F = 0.5
Delta= 1

#Length and Width of scattering region (the superconductor)
L = 10
W = 10


def make_system():
    lat = kwant.lattice.general([(sqrt(3)*0.5,0.5),(0,1)],[(0,0),(1/(2*sqrt(3)),1/2)],norbs = 2)
    a, b = lat.sublattices
    
    #Drawing the scattering region
    def scattering_region(pos):
        x,y = pos
        return abs(x) < L/2 and abs(y) < W/2
    #Drawing the lead region
    def lead_shape(pos):
        x, y = pos
        return abs(y) < W/2
    
    
    #Build the scattering region. Including Fermi energy, Superconducting gap, and nearest neighbour hopping.
    sys = kwant.Builder()
    sys[lat.shape(scattering_region,(0,0))] = E_F * tinyarray.array([[-1,0,0,0],[0,-1,0,0],[0,0,1,0],[0,0,0,1]])
    + Delta*tinyarray.array([[0,0,0,1],[0,0,-1,0],[0,-1,0,0],[1,0,0,0]]) 
    sys[lat.neighbors()] = t*tinyarray.array([[-1,0,0,0],[0,-1,0,0],[0,0,1,0],[0,0,0,1]])
    
    
    #######################################################################################################################
    
    #Building the lead. No superconductivity, only nearest neighbour hopping and Fermi energy.
    
    #What is the correct expression for the conservation law, and the particle hole symmetry?
    sym = kwant.TranslationalSymmetry(lat.vec((-2,1)))
    L_0 = kwant.Builder(sym,conservation_law=s_z)
    L_0[lat.shape(lead_shape,(0,0))] = E_F * tinyarray.array([[-1,0,0,0],[0,-1,0,0],[0,0,1,0],[0,0,0,1]])
    L_0[lat.neighbors()] = t*tinyarray.array([[-1,0,0,0],[0,-1,0,0],[0,0,1,0],[0,0,0,1]])
    
    ########################################################################################################################
    
    return sys, L_0


def plot_conductance(syst, energies):
    # Compute conductance
    data = []
    for energy in energies:
        smatrix = kwant.smatrix(syst, energy)
        # Conductance is N - R_ee + R_he
        data.append(smatrix.submatrix((0, 0), (0, 0)).shape[0] -
                    smatrix.transmission((0, 0), (0, 0)) +
                    smatrix.transmission((0, 1), (0, 0)))
    plt.figure()
    plt.plot(energies, data)
    plt.xlabel("energy [t]")
    plt.ylabel("conductance [e^2/h]")
    plt.show()

#Make my system and attach the leads
sys, L_0 = make_system()
sys.attach_lead(L_0)

#Plot the lattice with a lead attached to the scattering region
sysPlot = kwant.plot(sys,fig_size=(10,10))
sysPlot.savefig("System.pdf")

#Finalize my lead and system
L_0f = L_0.finalized()
sys = sys.finalized()

#Evaluate the band in the lead
bands = kwant.physics.Bands(L_0f)

#Plot the band in the lead
ks = np.linspace(-pi,pi,num=101)
energies=[bands(k) for k in ks]

f, ax = plt.subplots()
ax.plot(ks,energies)
ax.set_ylim(-1,1)
plt.show()

#Plot the conductance
plot_conductance(sys, energies=[0.01 * i for i in range(100)])