import math

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import OneHotEncoder
from tensorflow.keras import layers, losses
from tensorflow.keras import backend as K
from models.custom_layers import ExtractCentralMessage, OpticalChannel, DigitizationLayer, BitsToSymbols, SymbolsToBits
import itertools


class EndToEndAutoencoder(tf.keras.Model):
    def __init__(self,
                 cardinality,
                 samples_per_symbol,
                 messages_per_block,
                 channel,
                 bit_mapping=False):
        """
        The autoencoder that aims to find a encoding of the input messages. It should be noted that a "block" consists
        of multiple "messages" to introduce memory into the simulation as this is essential for modelling inter-symbol
        interference. The autoencoder architecture was heavily influenced by IEEE 8433895.

        :param cardinality: Number of different messages. Chosen such that each message encodes log_2(cardinality) bits
        :param samples_per_symbol: Number of samples per transmitted symbol
        :param messages_per_block: Total number of messages in transmission block
        :param channel: Channel Layer object. Must be a subclass of keras.layers.Layer with an implemented forward pass
        """
        super(EndToEndAutoencoder, self).__init__()

        # Labelled M in paper
        self.cardinality = cardinality
        self.bits_per_symbol = int(math.log(self.cardinality, 2))

        # Labelled n in paper
        self.samples_per_symbol = samples_per_symbol

        # Labelled N in paper
        if messages_per_block % 2 == 0:
            messages_per_block += 1
        self.messages_per_block = messages_per_block

        # Channel Model Layer
        if isinstance(channel, layers.Layer):
            self.channel = tf.keras.Sequential([
                layers.Flatten(),
                channel,
                ExtractCentralMessage(self.messages_per_block, self.samples_per_symbol)
            ], name="channel_model")
        else:
            raise TypeError("Channel must be a subclass of keras.layers.layer!")

        # Boolean identifying if bit mapping is to be learnt
        self.bit_mapping = bit_mapping

        # other parameters/metrics
        self.symbol_error_rate = None
        self.bit_error_rate = None
        self.snr = 20 * math.log(0.5 / channel.rx_stddev, 10)

        # Model Hyper-parameters
        leaky_relu_alpha = 0
        relu_clip_val = 1.0

        # Layer configuration for the case when bit mapping is to be learnt
        if self.bit_mapping:
            encoding_layers = [
                layers.Input(shape=(self.messages_per_block, self.bits_per_symbol)),
                BitsToSymbols(self.cardinality),
                layers.TimeDistributed(layers.Dense(2 * self.cardinality)),
                layers.TimeDistributed(layers.LeakyReLU(alpha=leaky_relu_alpha)),
                # layers.TimeDistributed(layers.Dense(2 * self.cardinality)),
                # layers.TimeDistributed(layers.LeakyReLU(alpha=leaky_relu_alpha)),
                # layers.TimeDistributed(layers.Dense(self.samples_per_symbol, activation='sigmoid')),
                layers.TimeDistributed(layers.Dense(self.samples_per_symbol)),
                layers.TimeDistributed(layers.ReLU(max_value=relu_clip_val))
            ]
            decoding_layers = [
                layers.Dense(2 * self.cardinality),
                layers.LeakyReLU(alpha=leaky_relu_alpha),
                # layers.Dense(2 * self.cardinality),
                # layers.LeakyReLU(alpha=0.01),
                layers.Dense(self.bits_per_symbol, activation='sigmoid')
            ]

        # layer configuration for the case when only symbol mapping is to be learnt
        else:
            encoding_layers = [
                layers.Input(shape=(self.messages_per_block, self.cardinality)),
                layers.TimeDistributed(layers.Dense(2 * self.cardinality)),
                layers.TimeDistributed(layers.LeakyReLU(alpha=leaky_relu_alpha)),
                layers.TimeDistributed(layers.Dense(2 * self.cardinality)),
                layers.TimeDistributed(layers.LeakyReLU(alpha=leaky_relu_alpha)),
                layers.TimeDistributed(layers.Dense(self.samples_per_symbol, activation='sigmoid')),
                # layers.TimeDistributed(layers.Dense(self.samples_per_symbol)),
                # layers.TimeDistributed(layers.ReLU(max_value=relu_clip_val))
            ]
            decoding_layers = [
                layers.Dense(2 * self.cardinality),
                layers.LeakyReLU(alpha=leaky_relu_alpha),
                layers.Dense(2 * self.cardinality),
                layers.LeakyReLU(alpha=leaky_relu_alpha),
                layers.Dense(self.cardinality, activation='softmax')
            ]

        # Encoding Neural Network
        self.encoder = tf.keras.Sequential([
            *encoding_layers
        ], name="encoding_model")

        # Decoding Neural Network
        self.decoder = tf.keras.Sequential([
            *decoding_layers
        ], name="decoding_model")

    def cost(self, y_true, y_pred):
        symbol_cost = losses.CategoricalCrossentropy()(y_true, y_pred)

        y_bits_true = SymbolsToBits(self.cardinality)(y_true)
        y_bits_pred = SymbolsToBits(self.cardinality)(y_pred)

        bit_cost = losses.BinaryCrossentropy()(y_bits_true, y_bits_pred)

        a = 1

        return symbol_cost + a * bit_cost

    def generate_random_inputs(self, num_of_blocks, return_vals=False):
        """
        A method that generates a list of one-hot encoded messages. This is utilized for generating the test/train data.

        :param num_of_blocks: Number of blocks to generate. A block contains multiple messages to be transmitted in
        consecutively to model ISI. The central message in a block is returned as the label for training.
        :param return_vals: If true, the raw decimal values of the input sequence will be returned
        """

        cat = [np.arange(self.cardinality)]
        enc = OneHotEncoder(handle_unknown='ignore', sparse=False, categories=cat)

        mid_idx = int((self.messages_per_block - 1) / 2)

        if self.bit_mapping:
            rand_int = np.random.randint(2, size=(num_of_blocks * self.messages_per_block * self.bits_per_symbol, 1))

            out = rand_int

            out_arr = np.reshape(out, (num_of_blocks, self.messages_per_block, self.bits_per_symbol))

            if return_vals:
                return out_arr, out_arr, out_arr[:, mid_idx, :]

        else:
            rand_int = np.random.randint(self.cardinality, size=(num_of_blocks * self.messages_per_block, 1))

            out = enc.fit_transform(rand_int)

            out_arr = np.reshape(out, (num_of_blocks, self.messages_per_block, self.cardinality))

            if return_vals:
                out_val = np.reshape(rand_int, (num_of_blocks, self.messages_per_block, 1))
                return out_val, out_arr, out_arr[:, mid_idx, :]

        return out_arr, out_arr[:, mid_idx, :]

    def train(self, num_of_blocks=1e6, epochs=1, batch_size=None, train_size=0.8, lr=1e-3):
        """
        Method to train the autoencoder. Further configuration to the loss function, optimizer etc. can be made in here.

        :param num_of_blocks: Number of blocks to generate for training. Analogous to the dataset size.
        :param batch_size: Number of samples to consider on each update iteration of the optimization algorithm
        :param train_size: Float less than 1 representing the proportion of the dataset to use for training
        :param lr: The learning rate of the optimizer. Defines how quickly the algorithm converges
        """
        X_train, y_train = self.generate_random_inputs(int(num_of_blocks * train_size))
        X_test, y_test = self.generate_random_inputs(int(num_of_blocks * (1 - train_size)))

        opt = tf.keras.optimizers.Adam(learning_rate=lr)

        # TODO: Investigate different optimizers (with different learning rates and other parameters)
        # SGD
        # RMSprop
        # Adam
        # Adadelta
        # Adagrad
        # Adamax
        # Nadam
        # Ftrl

        if self.bit_mapping:
            loss_fn = losses.BinaryCrossentropy()
        else:
            # loss_fn = losses.CategoricalCrossentropy()
            loss_fn = self.cost

        self.compile(optimizer=opt,
                     loss=loss_fn,
                     metrics=['accuracy'],
                     loss_weights=None,
                     weighted_metrics=None,
                     run_eagerly=False
                     )

        history = self.fit(x=X_train,
                           y=y_train,
                           batch_size=batch_size,
                           epochs=epochs,
                           shuffle=True,
                           validation_data=(X_test, y_test)
                           )

    def test(self, num_of_blocks=1e4, length_plot=False, plt_show=True):
        X_test, y_test = self.generate_random_inputs(int(num_of_blocks))

        y_out = self.call(X_test)

        y_pred = tf.argmax(y_out, axis=1)
        y_true = tf.argmax(y_test, axis=1)

        self.symbol_error_rate = 1 - accuracy_score(y_true, y_pred)

        bits_pred = SymbolsToBits(self.cardinality)(tf.one_hot(y_pred, self.cardinality)).numpy().flatten()
        bits_true = SymbolsToBits(self.cardinality)(y_test).numpy().flatten()

        self.bit_error_rate = 1 - accuracy_score(bits_true, bits_pred)

        if (length_plot):

            lengths = np.linspace(0, 70, 50)

            ber_l = []

            for l in lengths:
                tx_channel = OpticalChannel(fs=self.channel.layers[1].fs,
                                            num_of_samples=self.channel.layers[1].num_of_samples,
                                            dispersion_factor=self.channel.layers[1].dispersion_factor,
                                            fiber_length=l,
                                            lpf_cutoff=self.channel.layers[1].lpf_cutoff,
                                            rx_stddev=self.channel.layers[1].rx_stddev,
                                            sig_avg=self.channel.layers[1].sig_avg,
                                            enob=self.channel.layers[1].enob)

                test_channel = tf.keras.Sequential([
                    layers.Flatten(),
                    tx_channel,
                    ExtractCentralMessage(self.messages_per_block, self.samples_per_symbol)
                ], name="test channel (variable length)")

                X_test_l, y_test_l = self.generate_random_inputs(int(num_of_blocks))

                y_out_l = self.decoder(test_channel(self.encoder(X_test_l)))

                y_pred_l = tf.argmax(y_out_l, axis=1)
                # y_true_l = tf.argmax(y_test_l, axis=1)

                bits_pred_l = SymbolsToBits(self.cardinality)(tf.one_hot(y_pred_l, self.cardinality)).numpy().flatten()
                bits_true_l = SymbolsToBits(self.cardinality)(y_test_l).numpy().flatten()

                bit_error_rate_l = 1 - accuracy_score(bits_true_l, bits_pred_l)
                ber_l.append(bit_error_rate_l)

            plt.plot(lengths, ber_l)
            plt.yscale('log')
            if plt_show:
                plt.show()

        print("SYMBOL ERROR RATE: {}".format(self.symbol_error_rate))
        print("BIT ERROR RATE: {}".format(self.bit_error_rate))

        pass

    def view_encoder(self):
        '''
        A method that views the learnt encoder for each distint message. This is displayed as a plot with a subplot for
        each message/symbol.
        '''

        mid_idx = int((self.messages_per_block - 1) / 2)

        if self.bit_mapping:
            messages = np.zeros((self.cardinality, self.messages_per_block, self.bits_per_symbol))
            lst = [list(i) for i in itertools.product([0, 1], repeat=self.bits_per_symbol)]

            idx = 0
            for msg in messages:
                msg[mid_idx] = lst[idx]
                idx += 1

        else:
            # Generate inputs for encoder
            messages = np.zeros((self.cardinality, self.messages_per_block, self.cardinality))

            idx = 0
            for msg in messages:
                msg[mid_idx, idx] = 1
                idx += 1

        # Pass input through encoder and select middle messages
        encoded = self.encoder(messages)
        enc_messages = encoded[:, mid_idx, :]

        # Compute subplot grid layout
        i = 0
        while 2 ** i < self.cardinality ** 0.5:
            i += 1

        num_x = int(2 ** i)
        num_y = int(self.cardinality / num_x)

        # Plot all symbols
        fig, axs = plt.subplots(num_y, num_x, figsize=(2.5 * num_x, 2 * num_y))

        t = np.arange(self.samples_per_symbol)
        if isinstance(self.channel.layers[1], OpticalChannel):
            t = t / self.channel.layers[1].fs

        sym_idx = 0
        for y in range(num_y):
            for x in range(num_x):
                axs[y, x].plot(t, enc_messages[sym_idx].numpy().flatten(), 'x')
                axs[y, x].set_title('Symbol {}'.format(str(sym_idx)))
                sym_idx += 1

        for ax in axs.flat:
            ax.set(xlabel='Time', ylabel='Amplitude', ylim=(0, 1))

        for ax in axs.flat:
            ax.label_outer()

        plt.show()
        pass

    def view_sample_block(self):
        '''
        Generates a random string of input message and encodes them. In addition to this, the output is passed through
        digitization layer without any quantization noise for the low pass filtering.
        '''
        # Generate a random block of messages
        val, inp, _ = self.generate_random_inputs(num_of_blocks=1, return_vals=True)

        # Encode and flatten the messages
        enc = self.encoder(inp)
        flat_enc = layers.Flatten()(enc)
        chan_out = self.channel.layers[1](flat_enc)

        # Instantiate LPF layer
        lpf = DigitizationLayer(fs=self.channel.layers[1].fs,
                                num_of_samples=self.messages_per_block * self.samples_per_symbol,
                                sig_avg=0)

        # Apply LPF
        lpf_out = lpf(flat_enc)

        # Time axis
        t = np.arange(self.messages_per_block * self.samples_per_symbol)
        if isinstance(self.channel.layers[1], OpticalChannel):
            t = t / self.channel.layers[1].fs

        # Plot the concatenated symbols before and after LPF
        plt.figure(figsize=(2 * self.messages_per_block, 6))

        for i in range(1, self.messages_per_block):
            plt.axvline(x=t[i * self.samples_per_symbol], color='black')

        plt.plot(t, flat_enc.numpy().T, 'x')
        plt.plot(t, lpf_out.numpy().T)
        plt.plot(t, chan_out.numpy().flatten())
        plt.ylim((0, 1))
        plt.xlim((t.min(), t.max()))
        plt.title(str(val[0, :, 0]))
        plt.show()
        pass

    def call(self, inputs, training=None, mask=None):
        tx = self.encoder(inputs)
        rx = self.channel(tx)
        outputs = self.decoder(rx)
        return outputs


SAMPLING_FREQUENCY = 336e9
CARDINALITY = 32
SAMPLES_PER_SYMBOL = 32
MESSAGES_PER_BLOCK = 9
DISPERSION_FACTOR = -21.7 * 1e-24
FIBER_LENGTH = 50
FIBER_LENGTH_STDDEV = 5


if __name__ == '__main__':

    stddevs = [0, 1, 5, 10]
    legend = []

    for s in stddevs:
        optical_channel = OpticalChannel(fs=SAMPLING_FREQUENCY,
                                         num_of_samples=MESSAGES_PER_BLOCK * SAMPLES_PER_SYMBOL,
                                         dispersion_factor=DISPERSION_FACTOR,
                                         fiber_length=FIBER_LENGTH,
                                         fiber_length_stddev=s,
                                         lpf_cutoff=32e9,
                                         rx_stddev=0.01,
                                         sig_avg=0.5,
                                         enob=10)

        ae_model = EndToEndAutoencoder(cardinality=CARDINALITY,
                                       samples_per_symbol=SAMPLES_PER_SYMBOL,
                                       messages_per_block=MESSAGES_PER_BLOCK,
                                       channel=optical_channel,
                                       bit_mapping=False)

        print(ae_model.snr)

        ae_model.train(num_of_blocks=3e5, epochs=5)
        ae_model.test(length_plot=True, plt_show=False)
        # plt.legend(['{} +/- {}'.format(FIBER_LENGTH, s)])

        legend.append('{} +/- {}'.format(FIBER_LENGTH, s))

    plt.legend(legend)
    plt.show()
    plt.savefig('ber_vs_length.eps', format='eps')

    # ae_model.view_encoder()
    # ae_model.view_sample_block()
    # # ae_model.summary()
    # ae_model.encoder.summary()
    # ae_model.channel.summary()
    # ae_model.decoder.summary()
    pass