#!/usr/bin/env python2

from gnuradio import gr
from gnuradio import blocks
import pmt

class test_mux_stream(gr.top_block):
	def __init__(self):
		gr.top_block.__init__(self)
		self.tb = gr.top_block()

		src_data_a           = 10*[1,1]
		src_data_b           = 10*[2,2,2]
		#expected_result     = 10*[1,1,2,2,2,]

		src_a = blocks.vector_source_f(src_data_a)
		src_b = blocks.vector_source_f(src_data_b)
		tag_stream_a = blocks.stream_to_tagged_stream(gr.sizeof_float, 1, 2, 'packet_len_a')
		tag_stream_b = blocks.stream_to_tagged_stream(gr.sizeof_float, 1, 3, 'packet_len_b')

		mux = blocks.stream_mux(gr.sizeof_float, (2, 3))
		dst = blocks.vector_sink_f()

		self.tb.connect(src_a, tag_stream_a)
		self.tb.connect(src_b, tag_stream_b)
		self.tb.connect(tag_stream_a, (mux,0))
		self.tb.connect(tag_stream_b, (mux,1))
		self.tb.connect(mux, dst)
		self.tb.run()

		print(dst.data())

		# check the tags
		tags = dst.tags()
		for tag in tags:
			print("Key: " +pmt.to_python(tag.key)
				+" Value: " +str(pmt.to_python(tag.value))
				+" Offset: " +str(tag.offset) )

if __name__ == '__main__':
	try:
		test_mux_stream().run()
	except [[KeyboardInterrupt]]:
		pass