import math
import os
from multiprocessing import Pool

from sklearn.metrics import accuracy_score

from defs import Modulator, Demodulator, Channel
from models.basic import AWGNChannel
from misc import generate_random_bit_array
from models.optical_channel import OpticalChannel
import matplotlib.pyplot as plt
import numpy as np

CPU_COUNT = os.environ.get("CPU_COUNT", os.cpu_count())


def show_constellation(mod: Modulator, chan: Channel, demod: Demodulator, samples=1000):
    x = generate_random_bit_array(samples)
    x_mod = mod.forward(x)
    x_chan = chan.forward(x_mod)
    x_demod = demod.forward(x_chan)

    plt.plot(x_chan.rect_x[x], x_chan.rect_y[x], '+')
    plt.plot(x_chan.rect_x[:, 0][~x], x_chan.rect_y[:, 1][~x], '+')
    plt.plot(x_mod.rect_x[:, 0], x_mod.rect_y[:, 1], 'ro')
    axes = plt.gca()
    axes.set_xlim([-2, +2])
    axes.set_ylim([-2, +2])
    plt.grid()
    plt.show()
    print('Accuracy : ' + str())


def get_ber(mod, chan, demod, samples=1000):
    if samples % mod.N:
        samples += mod.N - (samples % mod.N)
    x = generate_random_bit_array(samples)
    x_mod = mod.forward(x)
    x_chan = chan.forward(x_mod)
    x_demod = demod.forward(x_chan)
    return 1 - accuracy_score(x, x_demod)


def get_AWGN_ber(mod, demod, samples=1000, start=-8., stop=5., steps=30):
    ber_x = np.linspace(start, stop, steps)
    ber_y = []
    for noise in ber_x:
        ber_y.append(get_ber(mod, AWGNChannel(noise), demod, samples=samples))
    return ber_x, ber_y


def __calc_ber(packed):
    # This function has to be outside get_Optical_ber in order to be pickled by pool
    mod, demod, noise, length, pulse_shape, samples = packed
    tx_channel = OpticalChannel(noise_level=noise, dispersion=-21.7, symbol_rate=10e9, sample_rate=400e9,
                                length=length, pulse_shape=pulse_shape, sqrt_out=True)
    return get_ber(mod, tx_channel, demod, samples=samples)


def get_Optical_ber(mod, demod, samples=1000, start=-8., stop=5., steps=30, length=100, pulse_shape='rect'):
    ber_x = np.linspace(start, stop, steps)
    ber_y = []
    print(f"Computing Optical BER.. 0/{len(ber_x)}", end='')
    with Pool(CPU_COUNT) as pool:
        packed_args = [(mod, demod, noise, length, pulse_shape, samples) for noise in ber_x]
        for i, ber in enumerate(pool.imap(__calc_ber, packed_args)):
            ber_y.append(ber)
            i += 1  # just offset by 1
            print(f"\rComputing Optical BER.. {i}/{len(ber_x)} ({i * 100 / len(ber_x):6.2f}%)", end='')
    print()
    return ber_x, ber_y


def get_SNR(mod, demod, ber_func=get_Optical_ber, samples=1000, start=-5, stop=15, **ber_kwargs):
    """
    SNR for optics and RF should be calculated the same, that is A^2
    Because P∝V² and P∝I²
    """
    x_mod = mod.forward(generate_random_bit_array(samples * mod.N))
    sig_power = [A ** 2 for A in x_mod.amplitude]
    av_sig_pow = np.mean(sig_power)
    av_sig_pow = math.log(av_sig_pow, 10)

    noise_start = -start + av_sig_pow
    noise_stop = -stop + av_sig_pow
    ber_x, ber_y = ber_func(mod, demod, samples, noise_start, noise_stop, **ber_kwargs)
    SNR = -ber_x + av_sig_pow
    return SNR, ber_y


def show_train_history(history, title="", save=None):
    from matplotlib import pyplot as plt

    epochs = range(1, len(history.epoch) + 1)
    if 'loss' in history.history:
        plt.plot(epochs, history.history['loss'], label='Training Loss')
    if 'accuracy' in history.history:
        plt.plot(epochs, history.history['accuracy'], label='Training Accuracy')
    if 'val_loss' in history.history:
        plt.plot(epochs, history.history['val_loss'], label='Validation Loss')
    if 'val_accuracy' in history.history:
        plt.plot(epochs, history.history['val_accuracy'], label='Validation Accuracy')
    plt.xlabel('Epochs')
    plt.ylabel('Loss/Accuracy' if 'accuracy' in history.history else 'Loss')
    plt.legend()
    plt.title(title)
    if save is not None:
        plt.savefig(save)
    plt.show()