import itertools
import math
import numpy as np
from models.custom_layers import OpticalChannel
from matplotlib import pyplot as plt

BANDWIDTH = 64e9

CARDINALITY = 16
DISPERSION_FACTOR = -21.7 * 1e-24
FIBER_LENGTH = 50
FIBER_LENGTH_STDDEV = 0
RX_STDDEV = 0.01
SIG_AVG = 0.5
ENOB = 10

# Number of OFDM symbols to simulate
OFDM_N = 50
# Number of OFDM subcarriers
K = 16
# length of the cyclic prefix: 25% of the block
CP = K // 4
# number of pilot carriers per OFDM block
P = 8
# The known value each pilot transmits
pilotValue = 3 + 3j
# DC offset used to ensure all values are positive
DC_OFFSET = 50
DC_OFFSET = np.asarray([DC_OFFSET])

SHOW_PLOTS = True

mapping_table = np.asarray([-3 - 3j, -3 - 1j, -3 + 3j, -3 + 1j,
                            -1 - 3j, -1 - 1j, -1 + 3j, -1 + 1j,
                            3 - 3j, 3 - 1j, 3 + 3j, 3 + 1j,
                            1 - 3j, 1 - 1j, 1 + 3j, 1 + 1j])

bits_per_symbol = int(math.log(CARDINALITY, 2))
bits_lst = [list(i) for i in itertools.product([0, 1], repeat=bits_per_symbol)]

# Set true to view plot of symbols as IQ plot
if SHOW_PLOTS:
    for idx, sym in enumerate(mapping_table):
        plt.plot(sym.real, sym.imag, 'bo')
        plt.text(sym.real, sym.imag + 0.2, str(bits_lst[idx])[1:-1], ha='center')

    plt.title('Symbols used in each subcarrier')
    plt.xlabel('I')
    plt.ylabel('Q')
    plt.xlim(-4, 4)
    plt.ylim(-4, 4)
    plt.show()


# All subcarriers used
allCarriers = np.arange(K)

# Identifying pilot carriers (and adding final subcarrier as a pilot for convenience)
pilotCarriers = allCarriers[::K // P]
pilotCarriers = np.hstack([pilotCarriers, np.array([allCarriers[-1]])])
P = P + 1

# Removing pilot carriers to obtain data carriers
dataCarriers = np.delete(allCarriers, pilotCarriers)

if SHOW_PLOTS:
    plt.plot(pilotCarriers, np.zeros_like(pilotCarriers), 'bo', label='pilot')
    plt.plot(dataCarriers, np.zeros_like(dataCarriers), 'ro', label='data')
    plt.show()

# Generate random symbols as integers and then map symbol values onto them
input_syms = np.random.randint(CARDINALITY, size=len(dataCarriers))
mapped_syms = mapping_table[input_syms]

# Generate the upper sideband of the OFDM symbol
enc_stream_upper_f = np.zeros(K, dtype=complex)
enc_stream_upper_f[pilotCarriers] = pilotValue
enc_stream_upper_f[dataCarriers] = mapped_syms

# Generate the lower sideband of the OFDM symbol
enc_stream_lower_f = np.conjugate(np.flip(enc_stream_upper_f))

# Combine the two sidebands with a DC offset to ensure values are always positive
enc_stream_f = np.concatenate((DC_OFFSET, enc_stream_upper_f, enc_stream_lower_f), axis=None)

if SHOW_PLOTS:
    f = np.fft.fftfreq(enc_stream_f.shape[0], d=1/BANDWIDTH)
    plt.plot(f, np.real(enc_stream_f), 'x')
    plt.plot(f, np.imag(enc_stream_f), 'x')
    # plt.xlim(-0.5e10, 0.5e10)
    plt.show()

# Take the inverse fourier transform
enc_stream_t = np.fft.ifft(enc_stream_f)

if SHOW_PLOTS:
    t = np.arange(len(enc_stream_t))*(1/BANDWIDTH)
    plt.plot(t, np.real(enc_stream_t))
    plt.plot(t, np.imag(enc_stream_t))
    plt.show()

# Take the real part to be transmitted via the channel
tx_stream = np.real(enc_stream_t)

optical_channel = OpticalChannel(fs=BANDWIDTH,
                                 num_of_samples=len(tx_stream),  # TODO: determine size of input to channel
                                 dispersion_factor=DISPERSION_FACTOR,
                                 fiber_length=FIBER_LENGTH,
                                 fiber_length_stddev=FIBER_LENGTH_STDDEV,
                                 lpf_cutoff=BANDWIDTH / 2,
                                 rx_stddev=RX_STDDEV,
                                 sig_avg=SIG_AVG,
                                 enob=ENOB)

# rx_stream = optical_channel(enc_stream_t)
#
# if SHOW_PLOTS:
#     t = np.arange(len(tx_stream))*(1/BANDWIDTH)
#     plt.plot(t, tx_stream)
#     plt.plot(t, rx_stream)
#     plt.show()