import unittest

import numpy as np
import tensorflow.compat.v1 as tf
tf.disable_eager_execution()
from tensorflow.python.client import device_lib


def rf(*s):
    return np.random.random(size=s)


class TestGaussShape(unittest.TestCase):
    """ Tests the GaussShape operator """
    def setUp(self):
        # Load the rime operation library
        from montblanc.impl.rime.tensorflow import load_tf_lib
        self.rime = load_tf_lib()
        # Obtain a list of GPU device specifications ['/gpu:0', '/gpu:1', ...]
        self.gpu_devs = [d.name for d in device_lib.list_local_devices()
                         if d.device_type == 'GPU']

    def test_gauss_shape(self):
        for FT in [np.float32, np.float64]:
            self._impl_test_gauss_shape(FT)

    def _impl_test_gauss_shape(self, FT):
        ngsrc, ntime, na, nchan = 10, 15, 7, 16
        nbl = na*(na-1)//2

        np_uvw = rf(ntime, na, 3).astype(FT)
        np_ant1, np_ant2 = [np.int32(x) for x in np.triu_indices(na, 1)]
        np_ant1, np_ant2 = (np.tile(np_ant1, ntime).reshape(ntime, nbl),
                            np.tile(np_ant2, ntime).reshape(ntime, nbl))
        np_frequency = np.linspace(1.4e9, 1.5e9, nchan).astype(FT)
        np_gauss_params = rf(3, ngsrc)*np.array([0.1, 0.1, 1.0])[:, np.newaxis]
        np_gauss_params = np_gauss_params.astype(FT)

        assert np_ant1.shape == (ntime, nbl), np_ant1.shape
        assert np_ant2.shape == (ntime, nbl), np_ant2.shape
        assert np_frequency.shape == (nchan,)

        np_args = [np_uvw, np_ant1, np_ant2, np_frequency, np_gauss_params]
        arg_names = ["uvw", "ant1", "ant2", "frequency", "gauss_params"]

        tf_args = [tf.Variable(v, name=n) for v, n in zip(np_args, arg_names)]

        def _pin_op(device, *tf_args):
            """ Pin operation to device """
            with tf.device(device):
                return self.rime.gauss_shape(*tf_args)

        # Pin operation to CPU
        cpu_op = _pin_op('/cpu:0', *tf_args)

        # Run the op on all GPUs
        gpu_ops = [_pin_op(d, *tf_args) for d in self.gpu_devs]

        # Initialise variables
        init_op = tf.compat.v1.global_variables_initializer()

        with tf.compat.v1.Session() as S:
            S.run(init_op)

            cpu_gauss_shape = S.run(cpu_op)

            for gpu_gauss_shape in S.run(gpu_ops):
                self.assertTrue(np.allclose(cpu_gauss_shape,
                                            gpu_gauss_shape))


if __name__ == "__main__":
    unittest.main()
