# -*- 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 acting on spin
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]])

#Specifying 2x2 Pauli matrices acting on particle-hole
tau_x = tinyarray.array([[0,1],[1,0]])
tau_y = tinyarray.array([[0,-1j],[1j,0]])
tau_z = tinyarray.array([[1,0],[0,-1]])
tau_0 = tinyarray.array([[1,0],[0,1]])

#Fix hopping, Fermi energi, superconducting gap, and voltage
t = 1
E_F = 0.001
Delta=0.001
V0 = 0.5*t

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


def make_system():
    lat = kwant.lattice.general([(sqrt(3)*0.5,0.5),(0,1)],[(0,0),(1/(2*sqrt(3)),1/2)],norbs = 4)
    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, Electric potential, Superconducting gap, and nearest neighbour hopping.
    sys = kwant.Builder()
    sys[lat.shape(scattering_region,(0,0))] = (E_F+V0) *-1* np.kron(tau_z,s_0)
    + Delta*-1*np.kron(tau_y,s_y) 
    sys[lat.neighbors()] = t*-1* np.kron(tau_z,s_0)
    
    
    #######################################################################################################################
    
    #Building the lead. No superconductivity, only nearest neighbour hopping and Fermi energy.
    
    sym = kwant.TranslationalSymmetry(lat.vec((-2,1)))
    L_0 = kwant.Builder(sym,conservation_law=np.kron(tau_z,s_0))
    L_0[lat.shape(lead_shape,(0,0))] = E_F*-1* np.kron(tau_z,s_0)
    L_0[lat.neighbors()] = t*-1* np.kron(tau_z,s_0)
    ########################################################################################################################
    
    return sys, L_0


    # Compute conductance
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.ylim(-0.5,2)
    plt.show()
    #print(data)

#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(-4*t,4*t)
plt.show()

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