import numpy as np
import math	
import cmath

import matplotlib.pyplot as plt

from bin16 import *
from fft import *

import os

# (integers)
# 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_Int(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
	
	#this is for keeping stuff as integers
	for i in range(len(b1)):
		b1W[i] = int(b1W[i])
		b2W[i] = int(b2W[i])
		b3W[i] = int(b3W[i])
		e1W[i] = int(e1W[i])
		e2W[i] = int(e2W[i])
	
	#remembering that fft_CT already divides the FFT by N	 	
 	B1 = fft_CT(b1W[offset:offset+N])
	B1 = B1[1:N/2+1]
 	B2 = fft_CT(b2W[offset:offset+N])
	B2 = B2[1:N/2+1]
 	B3 = fft_CT(b3W[offset:offset+N])
	B3 = B3[1:N/2+1]
 	E1 = fft_CT(e1W[offset:offset+N])
	E1 = E1[1:N/2+1]
 	E2 = fft_CT(e2W[offset:offset+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()]]
	
	#this we can keep as float
	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

# (integers)
# function that takes the time average of several SpectralMatrices
def Spectral_MatriceAvgTime_Int(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_Int(b1,b2,b3,e1,e2,fm,N,offset = 0)
	for k in range(1,NSM):
		Saux,w = Spectral_Matrice_Int(b1,b2,b3,e1,e2,fm,N,offset = k*NSM)
		S = SumSpectralMatrices(S,Saux)
	for i in range(5):
		for j in range(5):
				for k in range(len(S[i][j])):
					S[i][j][k] = complex(int(S[i][j][k].real/NSM),int(S[i][j][k].imag/NSM))
	return S, w
	
# (integers)
# "Power spectrum density of the Magnetic Field"
# it's being coded over 32 bits for now
def PSD_B_Int(S):
	PB = np.zeros(len(S[0][0]))
	PB = S[0][0] + S[1][1] + S[2][2]
	PB = abs(PB)
	return PB

# (integers)
# "Power spectrum density of the Electric Field"
# it's being coded over 32 bits for now
def PSD_E_Int(S):
	PE = np.zeros(len(S[3][3]))
	PE = S[3][3] + S[4][4]
	PE = abs(PE)
	return PE

# (integers)
# "Ellipticity of the electromagnetic wave"
# Coding over 4 bits, from 0 to 1
def ellipticity_Int(S):
	PB = PSD_B_Int(S)
	e = np.zeros(len(PB))
	e = 2.0/PB * np.sqrt(1.0*(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(e)):
		if math.isnan(e[i]):
			continue
		e[i] = quantN(e[i],0,1,4)
		e[i] = truncN(e[i],4)
	return e

# (integers)
# "Degree of polarization of the electromagnetic wave"
# Coding over 3 bits, from 0 to 1
def degree_polarization_Int(S):
	PB = PSD_B_Int(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.0*TrSCM2 - 1.0*TrSCM*TrSCM)/(2.0*TrSCM*TrSCM))
	for i in range(len(d)):
		if not(math.isnan(d[i])):
			d[i] = quantN(d[i],0,1,3)
			d[i] = truncN(d[i],3)
	return d
	
# (integers)
# "Normal wave vector"
# Coding over 8 bits for n1 and n2, 1 bit for the n3 (0 if n3 = +1, 1 if n3 = -1)
def normal_vector_Int(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)
	for i in range(len(n1)):
		if not(math.isnan(n1[i])):
			n1[i] = quantN(n1[i],0,1,8)
			n1[i] = truncN(n1[i],8)
		if not(math.isnan(n2[i])):
			n2[i] = quantN(n2[i],0,1,8)
			n2[i] = truncN(n2[i],8)
		if not(math.isnan(n3[i])):
			if n3[i] == -1:
				n3[i] = 1
			if n3[i] == +1:
				n3[i] = 0
	return [n1,n2,n3]

# (integers)
# "Z-component of the normalized Poynting flux"
# Coding over 8 bits
def poynting_vector_Int(S):
	return poynting
	
# (integers)
# "Phase velocity estimator"
# Coding over 8 bits
def phase_velocity_Int(S):
	return vp

# (integers)
# "Autocorrelation"
# it's being coded over 32 bits for now
def autocorrelation(S):
	return [S[0][0],S[1][1],S[2][2],S[3][3],S[4][4]]

# (integers)
# Normalized cross correlations
# Coding over 8 bits, for values between -1 and +1
def cross_correlations_Int(S):

	R12 = 1.0 * S[0][1].real/np.sqrt(1.0 * S[0][0] * S[1][1])
	R12 = quantN(R12,-1,+1,8)
	R12 = truncN(R12,8)
	R13 = 1.0 * S[0][2].real/np.sqrt(1.0 * S[0][0] * S[2][2])
	R13 = quantN(R13,-1,+1,8)
	R13 = truncN(R13,8)
	R23 = 1.0 * S[1][2].real/np.sqrt(1.0 * S[1][1] * S[2][2])
	R23 = quantN(R23,-1,+1,8)
	R23 = truncN(R23,8)
	R45 = 1.0 * S[3][4].real/np.sqrt(1.0 * S[3][3] * S[4][4])
	R45 = quantN(R45,-1,+1,8)
	R45 = truncN(R45,8)
	R14 = 1.0 * S[0][3].real/np.sqrt(1.0 * S[0][0] * S[3][3])
	R14 = quantN(R14,-1,+1,8)
	R14 = truncN(R14,8)
	R15 = 1.0 * S[0][4].real/np.sqrt(1.0 * S[0][0] * S[4][4])
	R15 = quantN(R15,-1,+1,8)
	R15 = truncN(R15,8)
	R24 = 1.0 * S[1][3].real/np.sqrt(1.0 * S[1][1] * S[3][3])
	R24 = quantN(R24,-1,+1,8)
	R24 = truncN(R24,8)
	R25 = 1.0 * S[1][4].real/np.sqrt(1.0 * S[1][1] * S[4][4])
	R25 = quantN(R25,-1,+1,8)
	R25 = truncN(R25,8)
	R34 = 1.0 * S[2][3].real/np.sqrt(1.0 * S[2][2] * S[3][3])
	R34 = quantN(R34,-1,+1,8)
	R34 = truncN(R34,8)
	R35 = 1.0 * S[2][4].real/np.sqrt(1.0 * S[2][2] * S[4][4])
	R35 = quantN(R35,-1,+1,8)
	R35 = truncN(R35,8)

	I12 = 1.0 * S[0][1].imag/np.sqrt(1.0 * S[0][0] * S[1][1])
	I12 = quantN(I12,-1,+1,8)
	I12 = truncN(I12,8)
	I13 = 1.0 * S[0][2].imag/np.sqrt(1.0 * S[0][0] * S[2][2])
	I13 = quantN(I13,-1,+1,8)
	I13 = truncN(I13,8)
	I23 = 1.0 * S[1][2].imag/np.sqrt(1.0 * S[1][1] * S[2][2])
	I23 = quantN(I23,-1,+1,8)
	I23 = truncN(I23,8)
	I45 = 1.0 * S[3][4].imag/np.sqrt(1.0 * S[3][3] * S[4][4])
	I45 = quantN(I45,-1,+1,8)
	I45 = truncN(I45,8)
	I14 = 1.0 * S[0][3].imag/np.sqrt(1.0 * S[0][0] * S[3][3])
	I14 = quantN(I14,-1,+1,8)
	I14 = truncN(I14,8)
	I15 = 1.0 * S[0][4].imag/np.sqrt(1.0 * S[0][0] * S[4][4])
	I15 = quantN(I15,-1,+1,8)
	I15 = truncN(I15,8)
	I24 = 1.0 * S[1][3].imag/np.sqrt(1.0 * S[1][1] * S[3][3])
	I24 = quantN(I24,-1,+1,8)
	I24 = truncN(I24,8)
	I25 = 1.0 * S[1][4].imag/np.sqrt(1.0 * S[1][1] * S[4][4])
	I25 = quantN(I25,-1,+1,8)
	I25 = truncN(I25,8)
	I34 = 1.0 * S[2][3].imag/np.sqrt(1.0 * S[2][2] * S[3][3])
	I34 = quantN(I34,-1,+1,8)
	I34 = truncN(I34,8)
	I35 = 1.0 * S[2][4].imag/np.sqrt(1.0 * S[2][2] * S[4][4])
	I35 = quantN(I35,-1,+1,8)
	I35 = truncN(I35,8)
	
	return [R12,R13,R23,R45,R14,R15,R24,R25,R34,R35],[I12,I13,I23,I45,I14,I15,I24,I25,I34,I35]

# (integers)
# this function takes a Spectral Matrice in the input and gives a list with all the associated Basic Parameters	
def BasicParameters_Int(S,w):
	PB = PSD_B_Int(S)
	PE = PSD_E_Int(S)
	d = degree_polarization_Int(S)
	e = ellipticity_Int(S)
	n = normal_vector_Int(S)
	return [PB,PE,d,e,n]

# (integers)
# this function saves plots in .pdf for each of the Basic Parameters associated with the Spectral Matrice in the input	
def BasicParameters_plot_Int(S,w,fm,mypath):
	
	if not os.path.isdir(mypath):
		os.makedirs(mypath)
	
	PB = PSD_B_Int(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_Int(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_Int(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_Int(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_Int(S)
	print "n1"
	print n1
	print "n2"
	print n2
	print "n3"
	print n3