from sklearn.preprocessing import OneHotEncoder
import numpy as np
from tensorflow.keras import layers

from end_to_end import load_model
from models.custom_layers import DigitizationLayer, OpticalChannel
from matplotlib import pyplot as plt
import math

# plot frequency spectrum of e2e model
def plot_e2e_spectrum(model_name=None):
    # Load pre-trained model
    ae_model, params = load_model(model_name=model_name)

    # Generate a list of random symbols (one hot encoded)
    cat = [np.arange(params["cardinality"])]
    enc = OneHotEncoder(handle_unknown='ignore', sparse=False, categories=cat)
    rand_int = np.random.randint(params["cardinality"], size=(10000, 1))
    out = enc.fit_transform(rand_int)

    # Encode the list of symbols using the trained encoder
    a = ae_model.encode_stream(out).flatten()

    # Pass the output of the encoder through LPF
    lpf = DigitizationLayer(fs=params["fs"],
                            num_of_samples=320000,
                            sig_avg=0)(a).numpy()

    # Plot the frequency spectrum of the signal
    freq = np.fft.fftfreq(lpf.shape[-1], d=1 / params["fs"])
    mul = np.exp(0.5j * params["dispersion_factor"] * params["fiber_length"] * np.power(2 * math.pi * freq, 2))

    a = np.fft.ifft(mul)
    a2 = np.power(a, 2)
    b = np.abs(np.fft.fft(a2))


    plt.plot(freq, np.fft.fft(lpf), 'x')
    plt.ylim((-500, 500))
    plt.xlim((-5e10, 5e10))
    plt.show()

    # plt.plot(freq, np.fft.fft(lpf), 'x')
    plt.plot(freq, b)
    plt.ylim((-500, 500))
    plt.xlim((-5e10, 5e10))
    plt.show()

def plot_e2e_encoded_output(model_name=None):
    # Load pre-trained model
    ae_model, params = load_model(model_name=model_name)

    # Generate a random block of messages
    val, inp, _ = ae_model.generate_random_inputs(num_of_blocks=1, return_vals=True)

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

    # Instantiate LPF layer
    lpf = DigitizationLayer(fs=params["fs"],
                            num_of_samples=params["messages_per_block"] * params["samples_per_symbol"],
                            sig_avg=0)

    # Apply LPF
    lpf_out = lpf(flat_enc)

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

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

    for i in range(1, params["messages_per_block"]):
        plt.axvline(x=t[i * params["samples_per_symbol"]], color='black')
    plt.axhline(y=0, color='black')
    plt.plot(t, flat_enc.numpy().T, 'x', label='output of encNN')
    plt.plot(t, lpf_out.numpy().T, label='optical field at tx')
    plt.plot(t, chan_out.numpy().flatten(), label='optical field at rx')
    plt.ylim((-0.1, 0.1))
    plt.xlim((t.min(), t.max()))
    plt.title(str(val[0, :, 0]))
    plt.legend(loc='upper right')
    plt.show()

if __name__ == '__main__':
    # plot_e2e_spectrum()
    plot_e2e_encoded_output()