"""
Layer for sampling OFDM MIMO channels from a set of SIMO vectors (based on Sionna's OFDM channel class)
"""
import sionna.ofdm
import tensorflow as tf
from sionna.channel import subcarrier_frequencies
from sionna.utils import expand_to_rank
from tensorflow.keras.layers import Layer
from sionna.channel.apply_ofdm_channel import ApplyOFDMChannel
import numpy as np

class customOFDMChannel(Layer):
    # pylint: disable=line-too-long
    r"""custom OFDM Channel
    Applies OFDM Channels generated by REMCOM
    Randomly selects UE possitions
    OFDM frequency responses preprocessed by MATLAB Script
    """

    def __init__(self, channel_set, resource_grid, add_awgn=True,
                normalize_channel=False, return_channel=False, chanIdxComb=None, randomSubSamplingChanIdx=False,
                dtype=tf.complex64, **kwargs):
        super().__init__(trainable=False, dtype=dtype, **kwargs)

        self._channels = tf.cast(channel_set[:, :, :resource_grid.fft_size], dtype=dtype)
        self._rg = resource_grid
        self._add_awgn = add_awgn
        self._normalize_channel = normalize_channel
        self._return_channel = return_channel

        self._chanIdxComb = tf.constant(chanIdxComb, dtype=tf.int32)
        self._randomSubSamplingChanIdx = randomSubSamplingChanIdx

        if normalize_channel:
            # Normalization is performed such that for each batch example and
            # link, the mean energy per resource grid is one.
            # Average over TX antennas, RX antennas, and subcarriers.
            c = tf.reduce_mean(tf.square(tf.abs(self._channels)), axis=(1, 2), keepdims=True)
            self._channels = self._channels / tf.complex(tf.sqrt(c), tf.constant(0., dtype.real_dtype))

    def build(self, input_shape):  # pylint: disable=unused-argument
        self._apply_channel = ApplyOFDMChannel(self._add_awgn, tf.as_dtype(self.dtype))

    def call(self, inputs):

        if self._add_awgn:
            x, no = inputs
        else:
            x = inputs

        batchsize = tf.shape(x)[0]
        n_tx = self._rg.num_tx
        num_rx_ant = tf.shape(self._channels)[1]

        # randomly select channel
        if self._chanIdxComb is None:
            # Limited randomnes, strategy (Matlab code)
            # itx = num_tx * randi((batch_size, num_tx), 0, int(num_channels/num_tx)-1) + (0:(num_tx-1))
            chan_batch_ue_idx = n_tx*tf.random.uniform((batchsize, n_tx), 0, int(tf.shape(self._channels)[0]/n_tx) - 1,
                                                       tf.int32)+tf.range(n_tx, dtype=tf.int32)
        else:
            if self._randomSubSamplingChanIdx:
                # randomly sample rows of channel index combinations
                # tensorflow seed is reset whenever the model graph is re-traced (so, all the models apply the same channel indices)
                numIdx = int(tf.shape(self._chanIdxComb)[0])
                rows = tf.reshape(tf.range(0, numIdx, dtype=tf.int64), [numIdx, 1])
                rows = tf.random.shuffle(rows)[:batchsize]
                # tf.print(rows[0])
                # [rows, _, _] = tf.random.uniform_candidate_sampler(cands, numIdx, batchsize, True, numIdx)
                chan_batch_ue_idx = tf.gather(self._chanIdxComb, tf.squeeze(rows), axis=0)
            else:
                chan_batch_ue_idx = self._chanIdxComb
        # select channels for all batch elements
        h_freq = tf.gather(self._channels, chan_batch_ue_idx, axis=0)
        # h_freq is [batch_size, num_tx, num_rx_ant, num_ofdm_channels]
        # expand to dimensions [batch_size, 1, num_tx, 1, num_rx_ant, 1, num_ofdm_channels]
        h_freq = tf.expand_dims(h_freq, axis=1)
        h_freq = tf.expand_dims(h_freq, axis=3)
        h_freq = tf.expand_dims(h_freq, axis=-2)
        # h_freq is [batch_size, 1, num_tx, 1, num_rx_ant, 1, num_ofdm_channels]
        h_freq = tf.transpose(h_freq, perm=[0, 1, 4, 2, 3, 5, 6])
        # block-fading channel same channels over same OFDM block (ie. for all OFDM symbols)
        h_freq = tf.tile(h_freq, [1, 1, 1, 1, 1, self._rg.num_ofdm_symbols, 1])
        # h_freq: [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_ofdm_symbols, fft_size]
        # reshape h_freq (force shape to be [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_ofdm_symbols, fft_size]
        chan_shape = tf.concat([[batchsize], [1], [num_rx_ant], [n_tx], [1], [self._rg.num_ofdm_symbols], [self._rg.fft_size]], 0)
        h_freq = tf.reshape(h_freq, chan_shape)

        if self._add_awgn:
            y = self._apply_channel([x, h_freq, no])
        else:
            y = self._apply_channel([x, h_freq])

        if self._return_channel:
            return y, h_freq
        else:
            return y
