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

saveFig = True

a = 1

t1 = 1.0
t2 = 1.5

eps = 0.0
  
model = kwant.lattice.general( [ ( a, 0 ), ( 0, 2*a ) ], [ ( 0, 0 ) ] )

subls = model.sublattices[0]

def transparency( tblock ):
    return ( 0 if tblock.shape[0] == 0 else
             np.sum( np.abs( tblock )**2 )/tblock.shape[0] )

def saveTab( fileName, x, y ):
    f = open( fileName, 'w' )
    for i in range( len( x ) ):
        f.write( '%f\t%f\n' % ( x[i], y[i] ) )
    f.close()

def make_system():
    
    def chain(pos):
        x, y = pos
        return abs(y) <= 1.99*a and ( -0.01 * a < x < 1.01 * a )

    def potential( site, _U ):
        x, y = site.pos
        return 0 if x < 0 else _U

    def onsite( site, _U=0 ):
        return eps + potential( site, _U )

    def lead0_onsite(site, _U):
        return eps

    def lead1_onsite(site, _U):
        return eps + _U

    syst = kwant.Builder()

    syst[model.shape(chain, (0, 0))] = onsite

    hopping1 = ( (-1, 0), subls, subls )
    hopping2 = ( (1, 0), subls, subls )

    syst[kwant.builder.HoppingKind( *hopping2 ) ] = t2

    sym0 = kwant.TranslationalSymmetry(model.vec((-1, 0)))
    def lead0_shape(pos):
        x, y = pos
        return ( abs(y) <= 1.99*a  )
        
    lead0 = kwant.Builder(sym0)
    lead0[model.shape(lead0_shape, (-a, 0))] = lead0_onsite
    lead0[kwant.builder.HoppingKind( *hopping1 ) ] = t1

    sym1 = kwant.TranslationalSymmetry(model.vec((1, 0)))
    def lead1_shape(pos):
        x, y = pos
        return ( abs(y) <= 1.99*a  )

    lead1 = kwant.Builder(sym1)
    lead1[model.shape(lead1_shape, ( 2 * a, 0))] = lead1_onsite
    lead1[kwant.builder.HoppingKind( *hopping2 ) ] = t2

    return syst, [lead0,lead1]


def plot_transparency( syst, energies, height = 0.0 ):
    data = []

    for i in range( len( energies ) ):
        smatrix = kwant.smatrix( syst, energies[i], args = [height] )
        #data.append( transparency( smatrix.submatrix( 1, 0 ) ) )
        data.append( smatrix.transmission( 1, 0 ) )
        
    plt.clf()
    plt.plot(energies, data, 'b', lw = 3)
    ax = plt.subplot(111)

    ax.tick_params(axis='both', which='major', labelsize=16)
    ax.tick_params(axis='both', which='minor', labelsize=12)
    plt.xlabel( 'Energy, a.u.', fontsize=20 )
    plt.ylabel( 'Transparency', fontsize=20 )
    plt.xlim( es[0], es[-1] )

    if saveFig:
        fn = 'testChain-kwant'
        saveTab( '%s.txt' % ( fn ), energies, data )
        plt.savefig( '%s.pdf' % ( fn ) )

    plt.show()

syst, leads = make_system()

kwant.plot(syst,site_size=0.18, site_lw=0.01, hop_lw=0.05)

for lead in leads:
    syst.attach_lead(lead)

kwant.plot(syst,site_size=0.18, site_lw=0.01, hop_lw=0.05)

syst = syst.finalized()

es = np.linspace( -1.99, 1.99, 500 )
plot_transparency( syst, es )