"""
Custom Keras Layers for general use
"""
import itertools

from tensorflow.keras import layers
import tensorflow as tf
import numpy as np


class AwgnChannel(layers.Layer):
    def __init__(self, rx_stddev=0.1, noise_dB=None, **kwargs):
        """
        :param rx_stddev: Standard deviation of receiver noise (due to e.g. TIA circuit)
        """
        super(AwgnChannel, self).__init__(**kwargs)
        if noise_dB is not None:
            # rx_stddev = np.sqrt(1 / (20 ** (noise_dB / 10.0)))
            rx_stddev = 10 ** (noise_dB / 10.0)
        self.noise_layer = layers.GaussianNoise(rx_stddev)

    def call(self, inputs, **kwargs):
        return self.noise_layer.call(inputs, training=True)


class ScaleAndOffset(layers.Layer):
    """
    Scales and offsets a tensor
    """

    def __init__(self, scale=1, offset=0, **kwargs):
        super(ScaleAndOffset, self).__init__(**kwargs)
        self.offset = offset
        self.scale = scale

    def call(self, inputs, **kwargs):
        return inputs * self.scale + self.offset


class BitsToSymbol(layers.Layer):
    def __init__(self, cardinality, **kwargs):
        super().__init__(**kwargs)
        self.cardinality = cardinality
        n = int(np.log(self.cardinality, 2))
        self.powers = tf.convert_to_tensor(
            np.power(2, np.linspace(n - 1, 0, n)).reshape(-1, 1),
            dtype=tf.float32
        )

    def call(self, inputs, **kwargs):
        idx = tf.cast(tf.tensordot(inputs, self.powers, axes=1), dtype=tf.int32)
        return tf.one_hot(idx, self.cardinality)


class SymbolToBits(layers.Layer):
    def __init__(self, cardinality, **kwargs):
        super().__init__(**kwargs)
        n = int(np.log(cardinality, 2))
        l = [list(i) for i in itertools.product([0, 1], repeat=n)]
        self.all_syms = tf.transpose(tf.convert_to_tensor(np.asarray(l), dtype=tf.float32))

    def call(self, inputs, **kwargs):
        return tf.matmul(self.all_syms, inputs)