#!/usr/bin/env python
#
# Copyright 2012,2013 Free Software Foundation, Inc.
#
# This file is part of GNU Radio
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
#


from gnuradio import gr, gr_unittest, blocks

from math import sqrt, atan2


class test_type_conversions(gr_unittest.TestCase):

    def setUp(self):
        self.tb = gr.top_block()

    def tearDown(self):
        self.tb = None

    def test_char_to_float_identity(self):
        src_data = (1, 2, 3, 4, 5)
        expected_data = [1.0, 2.0, 3.0, 4.0, 5.0]
        src = blocks.vector_source_b(src_data)
        op = blocks.char_to_float()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_char_to_float_scale(self):
        src_data = (1, 2, 3, 4, 5)
        expected_data = [0.5, 1.0, 1.5, 2.0, 2.5]
        src = blocks.vector_source_b(src_data)
        op = blocks.char_to_float(scale=2.0)
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_char_to_short(self):
        src_data = (1, 2, 3, 4, 5)
        expected_data = [256, 512, 768, 1024, 1280]
        src = blocks.vector_source_b(src_data)
        op = blocks.char_to_short()
        dst = blocks.vector_sink_s()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_complex_to_interleaved_char(self):
        src_data = (1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j, 9 + 10j)
        expected_data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_interleaved_char()
        dst = blocks.vector_sink_b()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_complex_to_interleaved_short(self):
        src_data = (1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j, 9 + 10j)
        expected_data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_interleaved_short()
        dst = blocks.vector_sink_s()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_complex_to_float_1(self):
        src_data = (1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j, 9 + 10j)
        expected_data = [1.0, 3.0, 5.0, 7.0, 9.0]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_float()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_complex_to_float_2(self):
        src_data = (1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j, 9 + 10j)
        expected_data1 = [1.0, 3.0, 5.0, 7.0, 9.0]
        expected_data2 = [2.0, 4.0, 6.0, 8.0, 10.0]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_float()
        dst1 = blocks.vector_sink_f()
        dst2 = blocks.vector_sink_f()
        self.tb.connect(src, op)
        self.tb.connect((op, 0), dst1)
        self.tb.connect((op, 1), dst2)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data1, dst1.data())
        self.assertFloatTuplesAlmostEqual(expected_data2, dst2.data())

    def test_complex_to_real(self):
        src_data = (1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j, 9 + 10j)
        expected_data = [1.0, 3.0, 5.0, 7.0, 9.0]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_real()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_complex_to_imag(self):
        src_data = (1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j, 9 + 10j)
        expected_data = [2.0, 4.0, 6.0, 8.0, 10.0]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_imag()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_complex_to_mag(self):
        src_data = (1 + 2j, 3 - 4j, 5 + 6j, 7 - 8j, -9 + 10j)
        expected_data = [sqrt(5), sqrt(25), sqrt(61), sqrt(113), sqrt(181)]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_mag()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data(), 5)

    def test_complex_to_mag_squared(self):
        src_data = (1 + 2j, 3 - 4j, 5 + 6j, 7 - 8j, -9 + 10j)
        expected_data = [5.0, 25.0, 61.0, 113.0, 181.0]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_mag_squared()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_complex_to_arg(self):
        src_data = (1 + 2j, 3 - 4j, 5 + 6j, 7 - 8j, -9 + 10j)
        expected_data = [atan2(2, 1), atan2(-4, 3),
                         atan2(6, 5), atan2(-8, 7), atan2(10, -9)]
        src = blocks.vector_source_c(src_data)
        op = blocks.complex_to_arg()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data(), 2)

    def test_float_to_char_identity(self):
        src_data = (1.0, 2.0, 3.0, 4.0, 5.0)
        expected_data = [1, 2, 3, 4, 5]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_char()
        dst = blocks.vector_sink_b()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_float_to_char_scale(self):
        src_data = (1.0, 2.0, 3.0, 4.0, 5.0)
        expected_data = [5, 10, 15, 20, 25]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_char(1, 5)
        dst = blocks.vector_sink_b()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_float_to_complex_1(self):
        src_data = (1.0, 3.0, 5.0, 7.0, 9.0)
        expected_data = [1 + 0j, 3 + 0j, 5 + 0j, 7 + 0j, 9 + 0j]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_complex()
        dst = blocks.vector_sink_c()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_float_to_complex_2(self):
        src1_data = (1.0, 3.0, 5.0, 7.0, 9.0)
        src2_data = (2.0, 4.0, 6.0, 8.0, 10.0)
        expected_data = [1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j, 9 + 10j]
        src1 = blocks.vector_source_f(src1_data)
        src2 = blocks.vector_source_f(src2_data)
        op = blocks.float_to_complex()
        dst = blocks.vector_sink_c()
        self.tb.connect(src1, (op, 0))
        self.tb.connect(src2, (op, 1))
        self.tb.connect(op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_float_to_int_identity(self):
        src_data = (1.0, 2.0, 3.0, 4.0, 5.0, float.fromhex('0x1.0p31'))
        expected_data = [1, 2, 3, 4, 5, 2147483647]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_int()
        dst = blocks.vector_sink_i()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_float_to_int_scale(self):
        src_data = (1.0, 2.0, 3.0, 4.0, 5.0)
        expected_data = [5, 10, 15, 20, 25]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_int(1, 5)
        dst = blocks.vector_sink_i()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_float_to_short_identity(self):
        src_data = (1.0, 2.0, 3.0, 4.0, 5.0)
        expected_data = [1, 2, 3, 4, 5]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_short()
        dst = blocks.vector_sink_s()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_float_to_short_scale(self):
        src_data = (1.0, 2.0, 3.0, 4.0, 5.0)
        expected_data = [5, 10, 15, 20, 25]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_short(1, 5)
        dst = blocks.vector_sink_s()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_float_to_uchar(self):
        src_data = (1.0, -2.0, 3.0, -4.0, 256.0)
        expected_data = [1, 0, 3, 0, 255]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_uchar()
        dst = blocks.vector_sink_b()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_float_to_uchar_scale_bias(self):
        src_data = (1.0, -2.0, 3.0, -4.0, 256.0)
        scale = 2.0
        bias = 5.0
        expected_data = [7, 1, 11, 0, 255]
        src = blocks.vector_source_f(src_data)
        op = blocks.float_to_uchar(scale=scale, bias=bias)
        dst = blocks.vector_sink_b()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_int_to_float_identity(self):
        src_data = (1, 2, 3, 4, 5)
        expected_data = [1.0, 2.0, 3.0, 4.0, 5.0]
        src = blocks.vector_source_i(src_data)
        op = blocks.int_to_float()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_int_to_float_scale(self):
        src_data = (1, 2, 3, 4, 5)
        expected_data = [0.2, 0.4, 0.6, 0.8, 1.0]
        src = blocks.vector_source_i(src_data)
        op = blocks.int_to_float(1, 5)
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertFloatTuplesAlmostEqual(expected_data, dst.data())

    def test_interleaved_short_to_complex(self):
        src_data = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
        expected_data = [1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j, 9 + 10j]
        src = blocks.vector_source_s(src_data)
        op = blocks.interleaved_short_to_complex()
        dst = blocks.vector_sink_c()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_short_to_char(self):
        src_data = (256, 512, 768, 1024, 1280)
        expected_data = [1, 2, 3, 4, 5]
        src = blocks.vector_source_s(src_data)
        op = blocks.short_to_char()
        dst = blocks.vector_sink_b()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_short_to_float_identity(self):
        src_data = (1, 2, 3, 4, 5)
        expected_data = [1.0, 2.0, 3.0, 4.0, 5.0]
        src = blocks.vector_source_s(src_data)
        op = blocks.short_to_float()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_short_to_float_scale(self):
        src_data = (5, 10, 15, 20, 25)
        expected_data = [1.0, 2.0, 3.0, 4.0, 5.0]
        src = blocks.vector_source_s(src_data)
        op = blocks.short_to_float(1, 5)
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())

    def test_uchar_to_float(self):
        src_data = (1, 2, 3, 4, 5)
        expected_data = [1.0, 2.0, 3.0, 4.0, 5.0]
        src = blocks.vector_source_b(src_data)
        op = blocks.uchar_to_float()
        dst = blocks.vector_sink_f()
        self.tb.connect(src, op, dst)
        self.tb.run()
        self.assertEqual(expected_data, dst.data())


if __name__ == '__main__':
    gr_unittest.run(test_type_conversions)
