import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
from matplotlib import gridspec
from sympy import *
from functions import *

def Vparcut(x):
	return max(0.,0.4-0.001 * x**2);

hsq=3.81 # h^2/2m_e

a0= 70.	#width of potential well
d= 200.	#total width

V0=0.3 	#well's depth
m=0.3	#effective mass

a=-115. 	#starting point
b=115. 		#ending point
N=300		#number of blocks


def x(i,start,end,blocks):
	return float(start+i*(end-start)/blocks);

def blocked(potential, start, end, blocks):
	blocked=np.empty([blocks])
	for i in range(blocks):
		blocked[i]= float((potential(x(i+1,start,end,blocks))+potential(x(i,start,end,blocks)))/2)
	return blocked;

def transmission(potential, start, end, blocks):
	E=0.35
	m2=1
	m1=1
	V=blocked(potential, start, end, blocks)

	M=np.empty([blocks,2,2], dtype=np.complex)
	M[0,]=np.identity(2)

	k0=sp.sqrt(m/hsq*(E-V[0]))
	trans = np.empty([blocks])
	refl = np.empty([blocks])

	for i in range(blocks-1):
		k1=sp.sqrt(m/hsq*(E-V[i]))
		k2=sp.sqrt(m/hsq*(E-V[i+1]))
		C=(k1*m2+k2*m1)/(2*k1*m2) +0j
		D=(k1*m2-k2*m1)/(2*k1*m2) +0j

		x0=x(i+1,start,end,blocks)
		M[i+1,]=np.matmul(M[i],([C*sp.exp(1j*(k2-k1)*x0), D*sp.exp(-1j*(k2+k1)*x0)],[D*sp.exp(1j*(k2+k1)*x0) , C*sp.exp(-1j*(k2-k1)*x0)]))

		if i>135 and i <155 :
			print('\n i='+str(i) +', x=' +str(x0) + ', k1=' + str(k1)+', k2=' + str(k2) )
			print(1/Abs(M[i,0,0])**2+ Abs(M[i,1,0]/M[i,0,0])**2)
			print([C*sp.exp(1j*(k2-k1)*x0), D*sp.exp(-1j*(k2+k1)*x0)],[D*sp.exp(1j*(k2+k1)*x0) , C*sp.exp(-1j*(k2-k1)*x0)])
			print(M[i,])

	print((1/Abs(M[blocks-1,0,0])**2+ Abs(M[blocks-1,1,0])**2/Abs(M[blocks-1,0,0])**2))

	for i in range(blocks):
		trans[i]=Abs(sp.sqrt(m/hsq*(E-V[i]))/(k0*M[i,0,0]**2))
		refl[i]=Abs(M[i,1,0]/M[i,0,0])**2


	print('transmission= ' +str(trans[blocks-1]))
	print('reflection= ' +str(refl[blocks-1]))

	grid = np.linspace(start,end,blocks)
	gs= gridspec.GridSpec(2,1,height_ratios=[2,1])

	a0=plt.subplot(gs[0])
	a0.plot(grid, trans, label='transmission')
	a0.plot(grid, refl, label='reflection')
	a0.legend(loc='upper left')
	#plt.title('Propability density')
	plt.xlabel('x [' + u"\u212B" +"]")
	#plt.ylabel(u"\u03A8"+' ^2')
	a0.xaxis.set_label_coords(1.05, -0.025)

	a1=plt.subplot(gs[1])
	plt.plot(grid, V)
	plt.title('Potential')
	plt.ylabel('V(x)')
	plt.xlabel('x [' + u"\u212B" +"]")
	a1.xaxis.set_label_coords(1.05, -0.025)

	plt.subplots_adjust(hspace=0.32)
	plt.savefig(potential.__name__+'.png', dpi=300)
	plt.close()
	return;

transmission(Vparcut,a,b,N)
#transmission(Vsq,a,b,N)