#
# SPDX-FileCopyrightText: Copyright (c) 2021-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
"""Layer for Convolutional Code Encoding."""

import tensorflow as tf
from tensorflow.keras.layers import Layer
from sionna.fec.utils import bin2int_tf, int2bin_tf
from sionna.fec.conv.utils import polynomial_selector, Trellis

class ConvEncoder(Layer):
    # pylint: disable=line-too-long
    r"""ConvEncoder(gen_poly=None, rate= 1/2, constraint_length=3, rsc=False, terminate=False, output_dtype=tf.float32, **kwargs)

    Encodes an information binary tensor to a convolutional codeword. Currently,
    only generator polynomials for codes of rate=1/n for n=2,3,4,... are allowed.

    The class inherits from the Keras layer class and can be used as layer in a
    Keras model.

    Parameters
    ----------
        gen_poly: tuple
            Sequence of strings with each string being a 0,1 sequence. If
            `None`, ``rate`` and ``constraint_length`` must be provided.

        rate: float
            Valid values are 1/3 and 0.5. Only required if ``gen_poly`` is
            `None`.

        constraint_length: int
            Valid values are between 3 and 8 inclusive. Only required if
            ``gen_poly`` is `None`.

        rsc: boolean
            Boolean flag indicating whether the Trellis generated is recursive
            systematic or not. If `True`, the encoder is recursive-systematic.
            In this case first polynomial in ``gen_poly`` is used as the
            feedback polynomial. Defaults to `False`.

        terminate: boolean
            Encoder is terminated to all zero state if `True`.
            If terminated, the `true` rate of the code is slightly lower than
            ``rate``.

        output_dtype: tf.DType
            Defaults to `tf.float32`. Defines the output datatype of the layer.

    Input
    -----
        inputs : [...,k], tf.float32
            2+D tensor containing the information bits where `k` is the
            information length

    Output
    ------
        : [...,k/rate], tf.float32
            2+D tensor containing the encoded codeword for the given input
            information tensor where `rate` is
            :math:`\frac{1}{\textrm{len}\left(\textrm{gen_poly}\right)}`
            (if ``gen_poly`` is provided).

    Note
    ----
        The generator polynomials from [Moon]_ are available for various
        rate and constraint lengths. To select them, use the ``rate`` and
        ``constraint_length`` arguments.

        In addition, polynomials for any non-recursive convolutional encoder
        can be given as input via ``gen_poly`` argument. Currently, only
        polynomials with rate=1/n are supported. When the ``gen_poly`` argument
        is given, the ``rate`` and ``constraint_length`` arguments are ignored.

        Various notations are used in the literature to represent the generator
        polynomials for convolutional codes. In [Moon]_, the octal digits
        format is primarily used. In the octal format, the generator polynomial
        `10011` corresponds to 46. Another widely used format
        is decimal notation with MSB. In this notation, polynomial `10011`
        corresponds to 19. For simplicity, the
        :class:`~sionna.fec.conv.encoding.ConvEncoder` only accepts the bit
        format i.e. `10011` as ``gen_poly`` argument.

        Also note that ``constraint_length`` and ``memory`` are two different
        terms often used to denote the strength of a convolutional code. In this
        sub-package, we use ``constraint_length``. For example, the
        polynomial `10011` has a ``constraint_length`` of 5, however its
        ``memory`` is only 4.

        When ``terminate`` is `True`, the true rate of the convolutional
        code is slightly lower than ``rate``. It equals
        :math:`\frac{r*k}{k+\mu}` where `r` denotes ``rate`` and
        :math:`\mu` is ``constraint_length`` - 1. For example when
        ``terminate`` is `True`, ``k=100``,
        :math:`\mu=4` and ``rate`` =0.5, true rate equals
        :math:`\frac{0.5*100}{104}=0.481`.
    """

    def __init__(self,
                 gen_poly=None,
                 rate=1/2,
                 constraint_length=3,
                 rsc=False,
                 terminate=False,
                 output_dtype=tf.float32,
                 **kwargs):

        super().__init__(**kwargs)

        if gen_poly is not None:
            assert all(isinstance(poly, str) for poly in gen_poly), \
                "Each element of gen_poly must be a string."
            assert all(len(poly)==len(gen_poly[0]) for poly in gen_poly), \
                "Each polynomial must be of same length."
            assert all(all(
                char in ['0','1'] for char in poly) for poly in gen_poly),\
                "Each Polynomial must be a string of 0/1 s."
            self._gen_poly = gen_poly
        else:
            valid_rates = (1/2, 1/3)
            valid_constraint_length = (3, 4, 5, 6, 7, 8)

            assert constraint_length in valid_constraint_length, \
                "Constraint length must be between 3 and 8."
            assert rate in valid_rates, \
                "Rate must be 1/3 or 1/2."
            self._gen_poly = polynomial_selector(rate, constraint_length)

        self._rsc = rsc
        self._terminate = terminate

        self._coderate_desired = 1/len(self.gen_poly)
        # Differ when terminate is True
        self._coderate = self._coderate_desired

        self._trellis = Trellis(self.gen_poly,rsc=self._rsc)
        self._mu = self.trellis._mu

        # conv_k denotes number of input bit streams.
        # Only 1 allowed in current implementation
        self._conv_k = self._trellis.conv_k

        # conv_n denotes number of output bits for conv_k input bits
        self._conv_n = self._trellis.conv_n

        self._ni = 2**self._conv_k
        self._no  = 2**self._conv_n
        self._ns = self._trellis.ns

        self._k = None
        self._n = None
        self.output_dtype = output_dtype

    #########################################
    # Public methods and properties
    #########################################

    @property
    def gen_poly(self):
        """Generator polynomial used by the encoder"""
        return self._gen_poly

    @property
    def coderate(self):
        """Rate of the code used in the encoder"""
        if self.terminate and self._k is None:
            print("Note that, due to termination, the true coderate is lower "\
                  "than the returned design rate. "\
                  "The exact true rate is dependent on the value of k and "\
                  "hence cannot be computed before the first call().")
        elif self.terminate and self._k is not None:
            term_factor = self._k/(self._k + self._mu)
            self._coderate = self._coderate_desired*term_factor
        return self._coderate

    @property
    def trellis(self):
        """Trellis object used during encoding"""
        return self._trellis

    @property
    def terminate(self):
        """Indicates if the convolutional encoder is terminated"""
        return self._terminate

    @property
    def k(self):
        """Number of information bits per codeword"""
        if self._k is None:
            print("Note: The value of k cannot be computed before the first " \
                  "call().")
        return self._k

    @property
    def n(self):
        """Number of codeword bits"""
        if self._n is None:
            print("Note: The value of n cannot be computed before the first " \
                  "call().")
        return self._n

    #########################
    # Keras layer functions
    #########################

    def build(self, input_shape):
        """Build layer and check dimensions.

        Args:
            input_shape: shape of input tensor (...,k).
        """
        self._k = input_shape[-1]
        self._n = int(self._k/self.coderate)

        # num_syms denote number of encoding periods or state transitions.
        # different from _k when _conv_k > 1.
        self.num_syms = int(self._k//self._conv_k)

    def call(self, inputs):
        """Convolutional code encoding function.

        Args:
            inputs (tf.float32): Information tensor of shape `[...,k]`.

        Returns:
            `tf.float32`: Encoded codeword tensor of shape `[...,n]`.
        """
        tf.debugging.assert_greater(tf.rank(inputs), 1)

        if inputs.shape[-1] != self._k:
            self.build(inputs.shape)

        msg = tf.cast(inputs, tf.int32)
        output_shape = msg.get_shape().as_list()
        output_shape[0] = -1 # overwrite batch dim (can be none in keras)
        output_shape[-1] = self._n # assign n to the last dim

        msg_reshaped = tf.reshape(msg, [-1, self._k])
        term_syms = int(self._mu) if self._terminate else 0

        prev_st = tf.zeros([tf.shape(msg_reshaped)[0]], tf.int32)
        ta = tf.TensorArray(tf.int32, size=self.num_syms, dynamic_size=False)

        idx_offset = range(0, self._conv_k)
        for idx in tf.range(0, self._k, self._conv_k):
            msg_bits_idx = tf.gather(msg_reshaped,
                                     idx + idx_offset,
                                     axis=-1)

            msg_idx = bin2int_tf(msg_bits_idx)

            indices = tf.stack([prev_st, msg_idx], -1)
            new_st = tf.gather_nd(self._trellis.to_nodes, indices=indices)

            idx_syms = tf.gather_nd(self._trellis.op_mat,
                                    tf.stack([prev_st, new_st], -1))
            idx_bits = int2bin_tf(idx_syms, self._conv_n)
            ta = ta.write(idx//self._conv_k, idx_bits)
            prev_st = new_st
        cw = tf.concat(tf.unstack(ta.stack()), axis=1)

        ta_term = tf.TensorArray(tf.int32, size=term_syms, dynamic_size=False)
        # Termination
        if self._terminate:
            if self._rsc:
                fb_poly = tf.constant([int(x) for x in self.gen_poly[0][1:]])
                fb_poly_tiled = tf.tile(
                        tf.expand_dims(fb_poly,0),[tf.shape(prev_st)[0],1])

            for idx in tf.range(0, term_syms, self._conv_k):
                prev_st_bits = int2bin_tf(prev_st, self._mu)
                if self._rsc:
                    msg_idx = tf.math.reduce_sum(
                                        tf.multiply(fb_poly_tiled, prev_st_bits),-1)
                    msg_idx = tf.squeeze(int2bin_tf(msg_idx,1),-1)
                else:
                    msg_idx = tf.zeros((tf.shape(prev_st)[0],), dtype=tf.int32)

                indices = tf.stack([prev_st, msg_idx], -1)
                new_st = tf.gather_nd(self._trellis.to_nodes, indices=indices)
                idx_syms = tf.gather_nd(self._trellis.op_mat,
                                        tf.stack([prev_st, new_st], -1))
                idx_bits = int2bin_tf(idx_syms, self._conv_n)
                ta_term = ta_term.write(idx//self._conv_k, idx_bits)
                prev_st = new_st

            term_bits = tf.concat(tf.unstack(ta_term.stack()), axis=1)
            cw = tf.concat([cw, term_bits], axis=-1)

        cw = tf.cast(cw, self.output_dtype)
        cw_reshaped = tf.reshape(cw, output_shape)

        return cw_reshaped

