from pycpsw import *
import numpy as np

# save files to 
f_dir = '/afs/slac/g/lcls/package/cpsw/controlGUI/MAIN_TRUNK/config/cryo/data/dataTest001'

# setup CPSW heirarchy
h = Hub.loadYamlFile( '/afs/slac/u/re/ulegat/ProjDocs/LLRF/firmware/targets/demo/AmcCarrierDemoEth/images/AmcCarrierDemoEth_project.yaml/000TopLevel.yaml', 'NetIODev' )

# setup PATHs 
p = Path.create( h )
p_mmio = p.findByName('mmio')
p_Demo = p_mmio.findByName( 'AmcCarrierDemoEth' )
p_AmcCore = p_Demo.findByName( 'AmcCarrierCore' )

# Load YAML config
p_mmio.loadConfigFromYamlFile( '/afs/slac/g/lcls/package/cpsw/controlGUI/MAIN_TRUNK/config/cryo/defaults.yaml' )


# create stream interfaces
strm0 = Stream.create( p.findByName('Stream0') )
strm1 = Stream.create( p.findByName('Stream1') )
strm2 = Stream.create( p.findByName('Stream2') )
strm3 = Stream.create( p.findByName('Stream3') )
# strm4 = Stream.create( p.findByName('Stream4') )
# strm5 = Stream.create( p.findByName('Stream5') )
# strm6 = Stream.create( p.findByName('Stream6') )
# strm7 = Stream.create( p.findByName('Stream7') )

# command to trigger daq
trig_cmd = Command.create( p_Demo.findByName('TrigDaq') )
init0    = Command.create( p_AmcCore.findByName('AmcCarrierBsa/BsaWaveformEngine[0]/WaveformEngineBuffers/Initialize') )
init1    = Command.create( p_AmcCore.findByName('AmcCarrierBsa/BsaWaveformEngine[1]/WaveformEngineBuffers/Initialize') )


daq_size = 0x1000
#daq_size = 0x2000

# Set WaveformEngineBuffers parameters
endAddr = ScalVal.create( p_mmio.findByName('AmcCarrierDemoEth/AmcCarrierCore/AmcCarrierBsa/BsaWaveformEngine/WaveformEngineBuffers/EndAddr') )
startAddr = ScalVal.create( p_mmio.findByName('AmcCarrierDemoEth/AmcCarrierCore/AmcCarrierBsa/BsaWaveformEngine/WaveformEngineBuffers/StartAddr') )
sAddr = np.array([0x0, 0x2000000, 0x4000000, 0x6000000, 0x8000000, 0xA000000, 0xC000000, 0xE000000])
eAddr = sAddr + 4*daq_size

# Define/Set DaqMuxV2 parameters: 
DataBufferSize = ScalVal.create( p_mmio.findByName('AmcCarrierDemoEth/AmcCarrierDemoApp/DaqMuxV2Bay1/DataBufferSize') )
DecimationRateDiv = ScalVal.create( p_mmio.findByName('AmcCarrierDemoEth/AmcCarrierDemoApp/DaqMuxV2Bay1/DecimationRateDiv') )
FormatSign = ScalVal.create( p_mmio.findByName('AmcCarrierDemoEth/AmcCarrierDemoApp/DaqMuxV2Bay1/FormatSign') )
FormatDataWidth = ScalVal.create( p_mmio.findByName('AmcCarrierDemoEth/AmcCarrierDemoApp/DaqMuxV2Bay1/FormatDataWidth') )
FormatSignWidth = ScalVal.create( p_mmio.findByName('AmcCarrierDemoEth/AmcCarrierDemoApp/DaqMuxV2Bay1/FormatSignWidth') )
DecimationAveraging = ScalVal.create( p_mmio.findByName('AmcCarrierDemoEth/AmcCarrierDemoApp/DaqMuxV2Bay1/DecimationAveraging') )

DataBufferSize.setVal( daq_size )
DecimationRateDiv.setVal( 0x04 )
FormatSign.setVal( 0x1 )
DecimationAveraging.setVal( 0x1 )
FormatDataWidth.setVal( 0x1 )
FormatSignWidth.setVal( 0x0F )

# setup Bsa/DaqMux for appropriate buffer size
endAddr.setVal( eAddr )
startAddr.setVal( sAddr )



# initialize BSA engine
init0.execute()
init1.execute()

# get numpy buffer, will use to store stream data
buf = np.getbuffer(np.zeros(daq_size*4), 0, daq_size*4)

dt = np.dtype('uint16')

# trigger daq
trig_cmd.execute()

# read buffers, remove header (8 bytes)
strm0.read(buf, daq_size)
newFile = open( f_dir + '_Stream0.bin', 'wb')
newFile.write( buf )
newFile.close()
ch0 = np.frombuffer(buf, dt, -1, 8)
ch0 = ch0.astype( np.float64 )
if dt == np.dtype('uint16'):
	ch0 = ch0 - 2**15

strm1.read(buf, daq_size)
newFile = open( f_dir + '_Stream1.bin', 'wb')
newFile.write( buf )
newFile.close()
ch1 = np.frombuffer(buf, dt, -1, 8)
ch1 = ch1.astype( np.float64 )
if dt == np.dtype('uint16'):
	ch1 = ch1 - 2**15

strm2.read(buf, daq_size)
newFile = open( f_dir + '_Stream2.bin', 'wb')
newFile.write( buf )
newFile.close()
ch2 = np.frombuffer(buf, dt, -1, 8)
ch2 = ch2.astype( np.float64 )
if dt == np.dtype('uint16'):
	ch2 = ch2 - 2**15

strm3.read(buf, daq_size)
newFile = open( f_dir + '_Stream3.bin', 'wb')
newFile.write( buf )
newFile.close()
ch3 = np.frombuffer(buf, dt, -1, 8)
ch3 = ch3.astype( np.float64 )
if dt == np.dtype('uint16'):
	ch3 = ch3 - 2**15

# strm4.read(buf, daq_size)
# newFile = open( f_dir + '_Stream4.bin', 'wb')
# newFile.write( buf )
# newFile.close()
# ch4 = np.frombuffer(buf, dt, -1, 8)
# ch4 = ch4.astype( np.float64 )
# if dt == np.dtype('uint16'):
	# ch4 = ch4 - 2**15

# strm5.read(buf, daq_size)
# newFile = open( f_dir + '_Stream5.bin', 'wb')
# newFile.write( buf )
# newFile.close()
# ch5 = np.frombuffer(buf, dt, -1, 8)
# ch5 = ch5.astype( np.float64 )
# if dt == np.dtype('uint16'):
	# ch5 = ch5 - 2**15

# strm6.read(buf, daq_size)
# newFile = open( f_dir + '_Stream6.bin', 'wb')
# newFile.write( buf )
# newFile.close()
# ch6 = np.frombuffer(buf, dt, -1, 8)
# ch6 = ch6.astype( np.float64 )
# if dt == np.dtype('uint16'):
	# ch6 = ch6 - 2**15

# strm7.read(buf, daq_size)
# newFile = open( f_dir + '_Stream7.bin', 'wb')
# newFile.write( buf )
# newFile.close()
# ch7 = np.frombuffer(buf, dt, -1, 8)
# ch7 = ch7.astype( np.float64 )
# if dt == np.dtype('uint16'):
	# ch7 = ch7 - 2**15