import numpy as np
import math	
import cmath

import matplotlib.pyplot as plt

from bin16 import *
from fft import *

import os

###########################################################

# (float)
# it calculates the Spectral Matrice for the input signals, taking N points in consideration and an optional offset for the starting point
def Spectral_Matrice(b1,b2,b3,e1,e2,fm,N,offset = 0):
	
	#the input signals shall pass by a Hann window before being analysed with the Fourier Transform
	W = np.zeros(len(b1))
	W[offset:offset+N] = np.array([math.pow((math.sin(math.pi*i/(N-1))),2) for i in range(0,N)])
	b1W = b1 * W
	b2W = b2 * W
	b3W = b3 * W
	e1W = e1 * W
	e2W = e2 * W

	B1 = np.fft.fft(b1W[offset:offset+N])/N
	B1 = B1[1:N/2+1]
	B2 = np.fft.fft(b2W[offset:offset+N])/N	
	B2 = B2[1:N/2+1]
	B3 = np.fft.fft(b3W[offset:offset+N])/N
	B3 = B3[1:N/2+1]
	E1 = np.fft.fft(e1W[offset:offset+N])/N
	E1 = E1[1:N/2+1]
	E2 = np.fft.fft(e2W[offset:offset+N])/N
	E2 = E2[1:N/2+1]
	
	#indices are SM[i][j][w] where i = line, j = colomn, w = frequency considered
	SM = [[B1*B1.conjugate(), B1*B2.conjugate(), B1*B3.conjugate(), B1*E1.conjugate(), B1*E2.conjugate()],
		  [B2*B1.conjugate(), B2*B2.conjugate(), B2*B3.conjugate(), B2*E1.conjugate(), B2*E2.conjugate()],
		  [B3*B1.conjugate(), B3*B2.conjugate(), B3*B3.conjugate(), B3*E1.conjugate(), B3*E2.conjugate()],
		  [E1*B1.conjugate(), E1*B2.conjugate(), E1*B3.conjugate(), E1*E1.conjugate(), E1*E2.conjugate()],
		  [E2*B1.conjugate(), E2*B2.conjugate(), E2*B3.conjugate(), E2*E1.conjugate(), E2*E2.conjugate()]]
	
	w = 2*math.pi*fm*np.arange(1,N/2+1)/N

	return SM, w

# function that sums two SpectralMatrices
def SumSpectralMatrices(M,N):
	Z = [[],[],[],[],[]]
	for i in range(5):
		for j in range(5):
			Z[i].append(M[i][j] + N[i][j])
	return Z

# (float)
# function that takes the time average of several SpectralMatrices
def Spectral_MatriceAvgTime(b1,b2,b3,e1,e2,fm,N):
	TBP = 4.0 # for the three cases of f in the LFR, TBP is equal to 4
	NSM = int(fm/N * TBP) # this gives the number of Spectral Matrices we should take into account for the average
	S,w = Spectral_Matrice(b1,b2,b3,e1,e2,fm,N,offset = 0)
	for k in range(1,NSM):
		Saux,w = Spectral_Matrice(b1,b2,b3,e1,e2,fm,N,offset = k*NSM)
		S = SumSpectralMatrices(S,Saux)
	for i in range(5):
		for j in range(5):
			S[i][j] = S[i][j]/NSM
	return S, w

# (float)
# "Power spectrum density of the Magnetic Field"
def PSD_B(S):
	PB = np.zeros(len(S[0][0]))
	PB = S[0][0] + S[1][1] + S[2][2]
	PB = abs(PB)
	return PB

# (float)
# "Power spectrum density of the Electric Field"
def PSD_E(S):
	PE = np.zeros(len(S[3][3]))
	PE = S[3][3] + S[4][4]
	PE = abs(PE)
	return PE

# (float)
# "Ellipticity of the electromagnetic wave"
def ellipticity(S):
	PB = PSD_B(S)
	e = np.zeros(len(PB))
	e = 2.0/PB * np.sqrt(S[0][1].imag*S[0][1].imag + S[0][2].imag*S[0][2].imag + S[1][2].imag*S[1][2].imag)
	return e

# (float)
# "Degree of polarization of the electromagnetic wave"
def degree_polarization(S):
	PB = PSD_B(S)
	d = np.zeros(len(PB))
	TrSCM2 = np.zeros(len(PB))
	TrSCM = np.zeros(len(PB))
	TrSCM2 = abs(S[0][0]*S[0][0] + S[1][1]*S[1][1] + S[2][2]*S[2][2]) + 2 * (abs(S[0][1]*S[0][1]) + abs(S[0][2]*S[0][2]) + abs(S[1][2]*S[1][2]))
	TrSCM  = PB
	d = np.sqrt((3*TrSCM2 - TrSCM*TrSCM)/(2*TrSCM*TrSCM))
	return d

# (float)
# "Normal wave vector"
def normal_vector(S):
	n1 = np.zeros(len(S[0][0]))
	n2 = np.zeros(len(S[0][0]))
	n3 = np.zeros(len(S[0][0]))
	n1 = +1.0 * S[1][2].imag/np.sqrt(S[0][1].imag*S[0][1].imag + S[0][2].imag*S[0][2].imag + S[1][2].imag*S[1][2].imag)
	n2 = -1.0 * S[0][2].imag/np.sqrt(S[0][1].imag*S[0][1].imag + S[0][2].imag*S[0][2].imag + S[1][2].imag*S[1][2].imag)
	for i in range(len(n3)):
		n3[i] = math.copysign(1,S[0][1][i].imag)
	return [n1,n2,n3]

# (float)	
# this is a routine that takes a list (or vector) of 3 components for the n vector and gives back an example
# of orthonormal matrix. This matrix may be used for the transformation matrix that takes every vector and puts it in the SCM referencial
def Matrix_SCM(n):

	n = np.array(n)
	n = n/np.linalg.norm(n)
	
	v = np.array([np.random.random(),np.random.random(),np.random.random()])

	v1 = v - np.dot(v,n)/np.dot(n,n) * n
	v1 = v1/np.linalg.norm(v1)
	
	v2 = np.cross(n,v1)
	v2 = v2/np.linalg.norm(v2)
	
	M = np.matrix([v1,v2,n]).transpose()
	
	return M

# (float)	
# this is a routine that takes the values 'a' and 'b' and creates a B complex vector (phasor of magnetic field)
# he then gives us the values of amplitude and phase for our cosines functions b1(t), b2(t), b3(t) for the SCM referencial (transformation using the R matrix) 	
def MagneticComponents(a,b,w,R):
	B_PA = np.matrix([a,-b*1j,0]).transpose()
	B_SCM = R*B_PA
	return B_SCM

# (float)
# this function takes some input values and gives back the corresponding Basic Parameters for the three possible fm's 
def SpectralMatrice_Monochromatic(a,b,n,E_para,f,fs,l):
	#'a' and 'b' are parameters related to the Magnetic Field
	#'n' is the normal vector to be considered
	#'E_para' is the degree of polarization of the propagating medium
	#'f' is the frequency of oscillation of the wave
	#'l' is the wavelength of the wave
	#'fm' is the original sampling frequency
	
	time_step = 1.0/fs
	t_ini = 0
	t_fin = 8
	time_vec = np.arange(t_ini,t_fin,time_step)

	# we must give the vector n, which has the information about the direction of k	
	n = n/np.linalg.norm(n) #we ensure that n is an unitary vector
	k = 2.0 * math.pi / l

	# the matrix R gives us an example of transformation to the SCM referencial
	# it was determined by the use of the normalized vector n and then two random vectors who form an orthogonal basis with it
	R = Matrix_SCM(n)

	# now we transform the magnetic field to the SCM referencial
	B = MagneticComponents(a,b,2.0 * math.pi * f,R)

	# now we have the caracteristics for our functions b1(t), b2(t) and b3(t) who shall enter in the LFR
	b1 = cmath.polar(B[0,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(B[0,0])[1])
	b2 = cmath.polar(B[1,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(B[1,0])[1])
	b3 = cmath.polar(B[2,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(B[2,0])[1])

	# the component of E who's parallel to n is of our choice. It represents the influence of the presence of charges where the electromagnetic
	# wave is travelling. However, the component of E who's perpendicular to n is determined by B and n as follows

	E_orth = -1.0 * 2.0 * math.pi * f/k * np.cross(n,np.array(B.transpose()))
	E = E_orth + E_para
	E = E.transpose()

	e1 = cmath.polar(E[0,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(E[0,0])[1])
	e2 = cmath.polar(E[1,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(E[1,0])[1])
	
	N = 256
	
	fm1 = 24576
	step1 = fs/fm1
	S1, w1 = Spectral_MatriceAvgTime(b1[::step1],b2[::step1],b3[::step1],e1[::step1],e2[::step1],fm1,N)

	fm2 = 4096
	step2 = fs/fm2
	S2, w2 = Spectral_MatriceAvgTime(b1[::step2],b2[::step2],b3[::step2],e1[::step2],e2[::step2],fm2,N)
	
	fm3 = 256	
	step3 = fs/fm3	
	S3, w3 = Spectral_MatriceAvgTime(b1[::step3],b2[::step3],b3[::step3],e1[::step3],e2[::step3],fm3,N)
	
	question = raw_input("Do you wish to save the plots (Y/N): ")
	if question == 'Y':
		BasicParameters_plot(S1,w1,fm1,"fm_%s" % fm1)
		BasicParameters_plot(S2,w2,fm2,"fm_%s" % fm2)
		BasicParameters_plot(S3,w3,fm3,"fm_%s" % fm3)
	
	return [S1,S2,S3], [w1,w2,w3]

# (float)
# this function takes a Spectral Matrice in the input and gives a list with all the associated Basic Parameters	
def BasicParameters(S,w):
	PB = PSD_B(S)
	PE = PSD_E(S)
	d = degree_polarization(S)
	e = ellipticity(S)
	n = normal_vector(S)
	return [PB,PE,d,e,n]

# (float)
# this function saves plots in .pdf for each of the Basic Parameters associated with the Spectral Matrice in the input	
def BasicParameters_plot(S,w,fm,mypath):
	
	if not os.path.isdir(mypath):
		os.makedirs(mypath)
	
	PB = PSD_B(S)	
	plt.plot(w/(2*math.pi),PB)
	plt.title("Power spectrum density of the Magnetic Field (fm = %s Hz)" % fm)
	plt.xlabel("frequency [Hz]")
	plt.ylabel("PB(w)")
	plt.savefig('%s/PB.png' % mypath)
	plt.close()

	PE = PSD_E(S)
	plt.plot(w/(2*math.pi),PE)
	plt.title("Power spectrum density of the Electric Field (fm = %s Hz)" % fm)
	plt.xlabel("frequency [Hz]")
	plt.ylabel("PE(w)")
	plt.savefig('%s/PE.png' % mypath)
	plt.close()

	e  = ellipticity(S)
	plt.plot(w/(2*math.pi),e)
	plt.title("Ellipticity of the electromagnetic wave (fm = %s Hz)" % fm)
	plt.xlabel("frequency [Hz]")
	plt.ylabel("e(w)")
	plt.savefig('%s/e.png' % mypath)
	plt.close()

	d  = degree_polarization(S)
	plt.plot(w/(2*math.pi),d)
	plt.title("Degree of polarization of the electromagnetic wave (fm = %s Hz)" % fm)
	plt.xlabel("frequency [Hz]")
	plt.ylabel("d(w)")
	plt.savefig('%s/d.png' % mypath)
	plt.close()	
	
	[n1, n2, n3]  = normal_vector(S)
	print "n1: "
	print n1
	print "n2: "	
	print n2
	print "n3: "	
	print n3

# (float)
# this function takes some input values and gives back the corresponding Basic Parameters for the three possible fm's 
def SignalsWaveMonochromatic(a,b,n,E_para,f,fs,l):
	#'a' and 'b' are parameters related to the Magnetic Field
	#'n' is the normal vector to be considered
	#'E_para' is the degree of polarization of the propagating medium
	#'f' is the frequency of oscillation of the wave
	#'l' is the wavelength of the wave
	#'fm' is the original sampling frequency
	
	time_step = 1.0/fs
	t_ini = 0
	t_fin = 8
	time_vec = np.arange(t_ini,t_fin,time_step)

	# we must give the vector n, which has the information about the direction of k	
	n = n/np.linalg.norm(n) #we ensure that n is an unitary vector
	k = 2.0 * math.pi / l
	
	print "   'a' = ", a
	print "   'b' = ", b
	print "     n = ", n
	print "     l = ", l
	print "     f = ", f
	print "E_para = ", E_para
	print " "
	
	# the matrix R gives us an example of transformation to the SCM referencial
	# it was determined by the use of the normalized vector n and then two random vectors who form an orthogonal basis with it
	R = Matrix_SCM(n)

	# now we transform the magnetic field to the SCM referencial
	B = MagneticComponents(a,b,2.0 * math.pi * f,R)

	# now we have the caracteristics for our functions b1(t), b2(t) and b3(t) who shall enter in the LFR
	b1 = cmath.polar(B[0,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(B[0,0])[1])
	b2 = cmath.polar(B[1,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(B[1,0])[1])
	b3 = cmath.polar(B[2,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(B[2,0])[1])
	
	print "b1(t) = ", cmath.polar(B[0,0])[0], "cos(w*t + ", cmath.polar(B[0,0])[1], ")"
	print "b2(t) = ", cmath.polar(B[1,0])[0], "cos(w*t + ", cmath.polar(B[1,0])[1], ")"
	print "b3(t) = ", cmath.polar(B[2,0])[0], "cos(w*t + ", cmath.polar(B[2,0])[1], ")"

	# the component of E who's parallel to n is of our choice. It represents the influence of the presence of charges where the electromagnetic
	# wave is travelling. However, the component of E who's perpendicular to n is determined by B and n as follows

	E_orth = -1.0 * 2.0 * math.pi * f/k * np.cross(n,np.array(B.transpose()))
	E = E_orth + E_para
	E = E.transpose()

	e1 = cmath.polar(E[0,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(E[0,0])[1])
	e2 = cmath.polar(E[1,0])[0] * np.cos(2.0 * math.pi * f * time_vec + cmath.polar(E[1,0])[1])
	
	print "e1(t) = ", cmath.polar(E[0,0])[0], "cos(w*t + ", cmath.polar(E[0,0])[1], ")"
	print "e2(t) = ", cmath.polar(E[1,0])[0], "cos(w*t + ", cmath.polar(E[1,0])[1], ")"
	print "(the value of w is ", 2.0 * math.pi * f, ")"
	print " "
	
	return [b1,b2,b3,e1,e2]