import numpy as np
import matplotlib.pyplot as plt
from errsim import *


def label(d, pe, pb, n):
    if pb is None:
        pb = pe
        label = 'd={} pe={} n={} BSC'.format(d, pe, n)
    else:
        label = 'd={} pe={} n={} pb={}'.format(d, pe, n, pb)
    return label

def plot(pe, fpath=None):
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=plt.figaspect(1/2))

    r = np.arange(8, 65)

    pWL = jointpmf5(pe, pe, 128)
    ax.plot(r, r_vs_pegd(pWL, 3, r) , 'g--', lw=2, label=label(3, pe, None, 128))
    ax.plot(r, r_vs_pegd(pWL, 6, r) , 'g-', lw=2, label=label(6, pe, None, 128))

    pWL = jointpmf5(pe, .1, 128)
    ax.plot(r, r_vs_pegd(pWL, 3, r) , 'b--', lw=2, label=label(3, pe, .1, 128))
    ax.plot(r, r_vs_pegd(pWL, 6, r) , 'b-', lw=2, label=label(6, pe, .1, 128))

    pWL = jointpmf5(pe, .5, 128)
    ax.plot(r, r_vs_pegd(pWL, 3, r) , 'r--', lw=2, label=label(3, pe, .5, 128))
    ax.plot(r, r_vs_pegd(pWL, 6, r) , 'r-', lw=2, label=label(6, pe, .5, 128))

    ax.set_yscale('log')
    ax.set_xticks(r[::8])
    ax.set_xlim(r[0], r[-1])
    #ax.set_ylim(1e-30, 1e-1)
    ax.set_xlabel('Burst error correction capability, $r$')
    ax.set_ylabel('$P_{egd}$')
    ax.set_title('Probability of Exceeding Guarenteed Error Detection Capability')

    ax.legend(loc='lower right')
    ax.grid(True)

    #plt.tight_layout()
    if fpath:
        fig.savefig(fpath)
    plt.show()

plt.close('all')
plot(1e-15, 'plots/pegd-pe=1e15.png')
plot(1e-6, 'plots/pegd-pe=1e6.png')
