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]) # nazywanie zmiennej tak samo jak funkcje to zla praktyka
	for i in range(blocks):
		blocked[i] = float((potential(x(i+1, start, end, blocks)) + potential(x(i, start, end, blocks))) / 2.)# fix potencjalnego buga - moze rzutowac na inta
	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): # na pewno -1? to nie sugestia zmiany tylko upewnienia sie, i te wszystkie nizej roznania polecam sprawdzic z naciskiem na +1, m2 m1 itd, tak w ogole nie polecam nazywac tak zmiennych, bardzo latwo sie jebnac
		k1 = sp.sqrt(m / hsq * (E - V[i]))
		k2 = sp.sqrt(m / hsq * (E - V[i + 1]))
     # (k1 * m2 + k2 * m1) dalabym w inna zmienna bo sie powtarza, a moze to zle i blad?
		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)
    # potwora nizej to wgl rozbic na zmienne
		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)
      # to nizej tez rozbic na zmienne
			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)