from math import pi
import numpy as np
import matplotlib.pyplot as plt

I = complex( 0, 1 )

saveFig = True

t1 = 1
t2 = 1.5

def disp( _t, _E, _p ):
    return 2 * _t * np.cos( _p ) - _E

def vPhi( _t, _p ):
    return -2 * _t * np.sin( _p )

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 transmission( _E, _plotBands = False ):
    
    _p1 = np.arccos( _E/(2*t1) )
    _p2 = np.arccos( _E/(2*t2) )
    
    if _plotBands:
        _ps = np.linspace( - pi, pi, 100 )
        plt.plot( _ps, disp( t1, _E, _ps ) + _E, 'b', label = 'disp1' )
        plt.plot( _ps, disp( t2, _E, _ps ) + _E, 'g', label = 'disp2' )
        plt.plot( [_ps[0], _ps[-1]], [ _E, _E ], 'r' )
        plt.legend( loc = 0 )
        plt.show()
    _A = np.array( [
        [ t1*np.exp(I*_p1)-_E, t2 ],
        [ t2, t2*np.exp(I*_p2)-_E ] ] )
    _b = np.array( [ _E - t1*np.exp(-I*_p1),
                     -t2 ] )
    
    x = np.linalg.solve( _A, _b )
    
    return [ x[0], x[-1], _p1, _p2 ]

E = -1.5

td = transmission( E, True )

print( td )

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

print( '|rho|**2 + |tau|**2 * v2/v1 = %f' % ( abs(td[0])**2 + abs(td[1])**2 * abs(vPhi( t2, td[3] )/vPhi( t1, td[2] )) ) )

trs = np.zeros( es.shape[0] )
for i in range( es.shape[0] ):
    rho, tau, p1, p2 = transmission( es[i] )
    trs[i] = abs(tau)**2 * abs(vPhi( t2, p2 )/vPhi( t1, p2 ))

plt.plot( es, trs, '--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] )
#plt.ylim( 0.3 , 1.0 )

if saveFig:
    fn = 'testChain'
    saveTab( '%s.txt' % ( fn ), es, trs )
    plt.savefig( '%s.pdf' % ( fn ) )

plt.show()