import numpy as np
import matplotlib.pyplot as plt

from test_fft.register_addresses import *

import test_fft.read_flags as flg
import test_fft.data_generator as dtgn
import test_fft.print_custom as prnt
import test_fft.read_data as rdata

prnt.print_custom( '*' )
prnt.print_custom( '*' )
prnt.print_custom( '*' )
prnt.print_custom( '*' )
prnt.print_custom( '*' )

############
# BUILD DATA
nbSamples = 256
nbComponentsPerMatrix = 15
nbFrequencyBins = 128
re_0 = dtgn.sineWave( nbFrequencyBins, 10, 10 )
im_0 = dtgn.sineWave( nbFrequencyBins, 10, 10 )
dataToWrite_0 = dtgn.generateDataToWrite( nbFrequencyBins, re_0, im_0 )
re_1 = dtgn.sineWave( nbFrequencyBins, 10, 20 )
im_1 = dtgn.sineWave( nbFrequencyBins, 10, 20 )
dataToWrite_1 = dtgn.generateDataToWrite( nbFrequencyBins, re_1, im_1 )
re_2 = dtgn.sineWave( nbFrequencyBins, 10, 30 )
im_2 = dtgn.sineWave( nbFrequencyBins, 10, 30 )
dataToWrite_2 = dtgn.generateDataToWrite( nbFrequencyBins, re_2, im_2 )
re_3 = dtgn.sineWave( nbFrequencyBins, 10, 40 )
im_3 = dtgn.sineWave( nbFrequencyBins, 10, 40 )
dataToWrite_3 = dtgn.generateDataToWrite( nbFrequencyBins, re_3, im_3 )
re_4 = dtgn.sineWave( nbFrequencyBins, 10, 50 )
im_4 = dtgn.sineWave( nbFrequencyBins, 10, 50 )
dataToWrite_4 = dtgn.generateDataToWrite( nbFrequencyBins, re_4, im_4 )

re_0 = dtgn.step( nbFrequencyBins, 10 )
im_0 = dtgn.continuous( nbFrequencyBins, 10 )
dataToWrite_0 = dtgn.generateDataToWrite( nbFrequencyBins, re_0, im_0 )
re_1 = dtgn.step( nbFrequencyBins, 20 )
im_1 = dtgn.continuous( nbFrequencyBins, 20 )
dataToWrite_1 = dtgn.generateDataToWrite( nbFrequencyBins, re_1, im_1 )
re_2 = dtgn.step( nbFrequencyBins, 30 )
im_2 = dtgn.continuous( nbFrequencyBins, 30 )
dataToWrite_2 = dtgn.generateDataToWrite( nbFrequencyBins, re_2, im_2 )
re_3 = dtgn.step( nbFrequencyBins, 40 )
im_3 = dtgn.continuous( nbFrequencyBins, 40 )
dataToWrite_3 = dtgn.generateDataToWrite( nbFrequencyBins, re_3, im_3 )
re_4 = dtgn.step( nbFrequencyBins, 50 )
im_4 = dtgn.continuous( nbFrequencyBins, 50 )
dataToWrite_4 = dtgn.generateDataToWrite( nbFrequencyBins, re_4, im_4 )

diracPosition = 60
re_0 = dtgn.dirac( nbFrequencyBins, 10, diracPosition )
im_0 = dtgn.dirac( nbFrequencyBins, 10, diracPosition )
dataToWrite_0 = dtgn.generateDataToWrite( nbFrequencyBins, re_0, im_0 )
re_1 = dtgn.dirac( nbFrequencyBins, 20, diracPosition )
im_1 = dtgn.dirac( nbFrequencyBins, 20, diracPosition )
dataToWrite_1 = dtgn.generateDataToWrite( nbFrequencyBins, re_1, im_1 )
re_2 = dtgn.dirac( nbFrequencyBins, 30, diracPosition )
im_2 = dtgn.dirac( nbFrequencyBins, 30, diracPosition )
dataToWrite_2 = dtgn.generateDataToWrite( nbFrequencyBins, re_2, im_2 )
re_3 = dtgn.dirac( nbFrequencyBins, 40, diracPosition )
im_3 = dtgn.dirac( nbFrequencyBins, 40, diracPosition )
dataToWrite_3 = dtgn.generateDataToWrite( nbFrequencyBins, re_3, im_3 )
re_4 = dtgn.dirac( nbFrequencyBins, 50, diracPosition )
im_4 = dtgn.dirac( nbFrequencyBins, 50, diracPosition )
dataToWrite_4 = dtgn.generateDataToWrite( nbFrequencyBins, re_4, im_4 )

########################
# WRITE WAVEFORM IN FIFO
prnt.print_reg_sm_only_MEM_IN_CTRL( SM_ONLY_MEM_IN_SM_CTRL )
prnt.print_reg_sm_only_MEM_OUT_CTRL( SM_ONLY_MEM_OUT_CTRL )
prnt.print_custom( "1) write data in SM_MEM_IN: " + str(len(dataToWrite_0)) + " bins" )

for k in range(nbFrequencyBins):
	RMAPPlugin0.Write( SM_ONLY_MEM_IN_SM_0, [dataToWrite_0[k]] )
	RMAPPlugin0.Write( SM_ONLY_MEM_IN_SM_1, [dataToWrite_1[k]] )
	RMAPPlugin0.Write( SM_ONLY_MEM_IN_SM_2, [dataToWrite_2[k]] )
	RMAPPlugin0.Write( SM_ONLY_MEM_IN_SM_3, [dataToWrite_3[k]] )
	RMAPPlugin0.Write( SM_ONLY_MEM_IN_SM_4, [dataToWrite_4[k]] )
	RMAPPlugin0.Write( SM_ONLY_MEM_IN_SM_CTRL, [0xffffffe0]   )

prnt.print_custom( "1) data written in SM_MEM_IN" )
prnt.print_reg_sm_only_MEM_IN_CTRL( SM_ONLY_MEM_IN_SM_CTRL )
prnt.print_reg_sm_only_MEM_OUT_CTRL( SM_ONLY_MEM_OUT_CTRL )

################
# SM FIRST READ
prnt.print_custom( "======= SM FIRST READ" )

components = []

flg.wait_for_FIFO_0_or_1_Full( SM_ONLY_MEM_OUT_CTRL )

if flg.is_MEM_OUT_SM_Full_FIFO_0( SM_ONLY_MEM_OUT_CTRL ):
	address_MEM_OUT_SM = SM_ONLY_MEM_OUT_SM_1
	fifo_to_wait_for = 0
elif flg.is_MEM_OUT_SM_Full_FIFO_1( SM_ONLY_MEM_OUT_CTRL ):
	address_MEM_OUT_SM = SM_ONLY_MEM_OUT_SM_0
	fifo_to_wait_for = 1
else:
	prnt.print_custom("oups, FIFO_0 and FIFO_1 are not full, this is unexpected")

for component in range(nbComponentsPerMatrix):
	prnt.print_custom( "==== component = " + str( component ) 
		+ ", read @ " + hex(address_MEM_OUT_SM & 0xffffffff) )
	prnt.print_reg_sm_only_MEM_IN_CTRL( SM_ONLY_MEM_IN_SM_CTRL )
	prnt.print_reg_sm_only_MEM_OUT_CTRL( SM_ONLY_MEM_OUT_CTRL )
	currentComp_re = np.zeros( nbFrequencyBins )
	currentComp_im = np.zeros( nbFrequencyBins )
	if (component == 0) | (component == 12):
		flg.wait_for_FIFO_Full( fifo_to_wait_for, SM_ONLY_MEM_OUT_CTRL )
		currentComp_re = rdata.read_SM_ONLY_Re( nbFrequencyBins, address_MEM_OUT_SM )
		components.append( currentComp_re )
	elif (component == 5) | (component == 9):
		flg.wait_for_FIFO_Full( fifo_to_wait_for, SM_ONLY_MEM_OUT_CTRL )
		currentComp_re = rdata.read_SM_ONLY_Re( nbFrequencyBins, address_MEM_OUT_SM )
		components.append( currentComp_re )
	elif (component == 1) | (component == 3) | (component == 7):
		currentComp_re, currentComp_im = rdata.read_SM_ONLY_Re_Im( nbFrequencyBins, address_MEM_OUT_SM )
		components.append( currentComp_re )
		components.append( currentComp_im )
		flg.wait_for_FIFO_Full( fifo_to_wait_for, SM_ONLY_MEM_OUT_CTRL )
	elif (component == 2) | (component == 6) | (component == 10):
		currentComp_re, currentComp_im = rdata.read_SM_ONLY_Re_Im( nbFrequencyBins, address_MEM_OUT_SM )
		components.append( currentComp_re )
		components.append( currentComp_im )
		flg.wait_for_FIFO_Full( fifo_to_wait_for, SM_ONLY_MEM_OUT_CTRL )
	elif (component == 4) | (component == 8):
		currentComp_re, currentComp_im = rdata.read_SM_ONLY_Re_Im( nbFrequencyBins, address_MEM_OUT_SM )
		components.append( currentComp_re )
		components.append( currentComp_im )
	elif (component == 11) | (component == 13):
		currentComp_re, currentComp_im = rdata.read_SM_ONLY_Re_Im( nbFrequencyBins, address_MEM_OUT_SM )
		components.append( currentComp_re )
		components.append( currentComp_im )
	elif (component == 14):
		currentComp_re = rdata.read_SM_ONLY_Re( nbFrequencyBins, address_MEM_OUT_SM )
		components.append( currentComp_re )
	else:
		prnt.print_custom( "unexpected value for component" )
	# change FIFO
	if address_MEM_OUT_SM == SM_ONLY_MEM_OUT_SM_0:
		address_MEM_OUT_SM = SM_ONLY_MEM_OUT_SM_1
		fifo_to_wait_for = 0
	elif address_MEM_OUT_SM == SM_ONLY_MEM_OUT_SM_1:
		address_MEM_OUT_SM = SM_ONLY_MEM_OUT_SM_0
		fifo_to_wait_for = 1

prnt.print_custom( "======= READ COMPLETE" )

prnt.print_reg_sm_only_MEM_IN_CTRL( SM_ONLY_MEM_IN_SM_CTRL )
prnt.print_reg_sm_only_MEM_OUT_CTRL( SM_ONLY_MEM_OUT_CTRL )

# PLOT SM

plt.figure( 1 )
plt.subplot(231)
plt.plot(re_0, 'b')
plt.subplot(232)
plt.plot(re_1, 'g')
plt.plot(re_1, '.')
plt.subplot(233)
plt.plot(re_2, 'r')
plt.plot(re_2, '.')
plt.subplot(234)
plt.plot(re_3, 'c')
plt.subplot(235)
plt.plot(re_4, 'm')

plt.figure( 2 )

plt.subplot(231)
plt.plot(components[0], label='0')
plt.plot(components[1], '.', label='1')
plt.plot(components[2], label='2')
plt.plot(components[3], '.', label='3')
plt.legend(loc='upper right')

plt.subplot(232)
plt.plot(components[4], label='4')
plt.plot(components[5], '.', label='5')
plt.plot(components[6], label='6')
plt.plot(components[7], '.', label='7')
plt.legend(loc='upper right')

plt.subplot(233)
plt.plot(components[8], label='8')
plt.plot(components[9], label='9')
plt.plot(components[10], '.', label='10')
plt.plot(components[11], label='11')
plt.legend(loc='upper right')

plt.subplot(234)
plt.plot(components[12], label='12')
plt.plot(components[13], label='13')
plt.plot(components[14], label='14')
plt.plot(components[15], label='15')
plt.legend(loc='upper right')

plt.subplot(235)
plt.plot(components[16], label='16')
plt.plot(components[17], label='17')
plt.plot(components[18], label='18')
plt.plot(components[19], label='19')
plt.legend(loc='upper right')

plt.subplot(236)
plt.plot(components[20], label='20')
plt.plot(components[21], label='21')
plt.plot(components[22], label='22')
plt.plot(components[23], label='23')
plt.plot(components[24], label='24')
plt.legend(loc='upper right')

plt.show()