import matplotlib.pyplot as plt

import graphs
from models.basic import AWGNChannel, BPSKDemod, BPSKMod, BypassChannel, AlphabetMod, AlphabetDemod


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=2e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 5bit -25dB')

    # view_encoder(aenc.encoder, 5)
    # plt.plot(*get_AWGN_ber(AlphabetMod('32qam', 10e6), AlphabetDemod('32qam', 10e6), samples=12000, start=-15), '-',
    #          label='32-QAM')
    # show_constellation(AlphabetMod('32qam', 10e6), AWGNChannel(-1), AlphabetDemod('32qam', 10e6))
    # mod = AlphabetMod('32qam', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # pass

    # aenc = Autoencoder(5, -15)
    # aenc.train(samples=2e6)
    # 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')

    # for scheme in ['64qam', '32qam', '16qam', 'qpsk', '8psk']:
    #     plt.plot(*get_SNR(
    #         AlphabetMod(scheme, 10e6),
    #         AlphabetDemod(scheme, 10e6),
    #         samples=100e3,
    #         steps=40,
    #         start=-15
    #     ), '-', label=scheme.upper())
    # plt.yscale('log')
    # plt.grid()
    # plt.xlabel('SNR dB')
    # plt.ylabel('BER')
    # plt.legend()
    # plt.show()

    # for l in np.logspace(start=0, stop=3, num=6):
    #     plt.plot(*misc.get_SNR(
    #         AlphabetMod('4pam', 10e6),
    #         AlphabetDemod('4pam', 10e6),
    #         samples=2000,
    #         steps=200,
    #         start=-5,
    #         stop=20,
    #         length=l,
    #         pulse_shape='rcos'
    #     ), '-', label=(str(int(l))+'km'))
    #
    # plt.yscale('log')
    # # plt.gca().invert_xaxis()
    # plt.grid()
    # plt.xlabel('SNR dB')
    # # plt.ylabel('BER')
    # plt.title("BER against Fiber length")
    # plt.legend()
    # plt.show()

    for ps in ['rect', 'rcos', 'rrcos']:
        plt.plot(*graphs.get_SNR(
            AlphabetMod('4pam', 10e6),
            AlphabetDemod('4pam', 10e6),
            samples=30000,
            steps=100,
            start=-5,
            stop=20,
            length=1,
            pulse_shape=ps
        ), '-', label=ps)

    plt.yscale('log')
    plt.grid()
    plt.xlabel('SNR dB')
    plt.ylabel('BER')
    plt.title("BER for different pulse shapes")
    plt.legend()
    plt.show()

    pass