import matplotlib.pyplot as plt

import numpy as np
from sklearn.metrics import accuracy_score
from models import basic
from models.basic import AWGNChannel, BPSKDemod, BPSKMod, BypassChannel, AlphabetMod, AlphabetDemod
import misc
from models.autoencoder import Autoencoder


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

    x_mod_rect = misc.polar2rect(x_mod)
    x_chan_rect = misc.polar2rect(x_chan)
    plt.plot(x_chan_rect[:, 0][x], x_chan_rect[:, 1][x], '+')
    plt.plot(x_chan_rect[:, 0][~x], x_chan_rect[:, 1][~x], '+')
    plt.plot(x_mod_rect[:, 0], x_mod_rect[:, 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 = misc.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


if __name__ == '__main__':
    # show_constellation(BPSKMod(10e6), AWGNChannel(-1), BPSKDemod(10e6, 10e3))

    # get_ber(BPSKMod(10e6), AWGNChannel(-20), BPSKDemod(10e6, 10e3))
    # mod = MaryMod('8psk', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # mod = MaryMod('qpsk', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # mod = MaryMod('16qam', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # mod = MaryMod('64qam', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # aenc = Autoencoder(4, -25)
    # aenc.train(samples=5e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 4bit -25dB')

    aenc = Autoencoder(5, -25)
    aenc.train(samples=1e6)
    plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
             label='AE 5bit -25dB')

    aenc = Autoencoder(5, -15)
    aenc.train(samples=1e6)
    plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
             label='AE 5bit -15dB')

    aenc = Autoencoder(4, -25)
    aenc.train(samples=6e5)
    plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
             label='AE 4bit -20dB')

    aenc = Autoencoder(4, -15)
    aenc.train(samples=6e5)
    plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
             label='AE 4bit -15dB')

    aenc = Autoencoder(2, -20)
    aenc.train(samples=6e5)
    plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
             label='AE 2bit -20dB')

    aenc = Autoencoder(2, -15)
    aenc.train(samples=6e5)
    plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
             label='AE 2bit -15dB')

    # aenc = Autoencoder(4, -10)
    # aenc.train(samples=5e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 4bit -10dB')
    #
    # aenc = Autoencoder(4, -8)
    # aenc.train(samples=5e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 4bit -8dB')

    plt.plot(*get_AWGN_ber(AlphabetMod('64qam', 10e6), AlphabetDemod('64qam', 10e6), samples=12000, start=-15), '-', label='64-QAM')
    plt.plot(*get_AWGN_ber(AlphabetMod('16qam', 10e6), AlphabetDemod('16qam', 10e6), samples=12000, start=-15), '-', label='16-QAM')
    plt.plot(*get_AWGN_ber(AlphabetMod('qpsk', 10e6), AlphabetDemod('qpsk', 10e6), samples=12000, start=-15), '-', label='QPSK')
    plt.plot(*get_AWGN_ber(AlphabetMod('8psk', 10e6), AlphabetDemod('8psk', 10e6), samples=12000, start=-15), '-', label='8PSK')
    plt.plot(*get_AWGN_ber(BPSKMod(10e6), BPSKDemod(10e6, 10e3), samples=12000), '-', label='BPSK')
    plt.yscale('log')
    plt.gca().invert_xaxis()
    plt.grid()
    plt.xlabel('Noise dB')
    plt.ylabel('BER')
    plt.legend()
    plt.show()

    pass