import numpy as np
import math as math
from pylab import * 
import py.ldpc as ldpc
import matplotlib.pyplot as plt
import time
import csv
from bitarray import bitarray
import amp_exit as ae


# Code taken from Adam's python tutorial. 

try:
    from pyfht import fht_inplace
except ImportError:
    import warnings
    warnings.warn("Using very slow Python version of fht, please install pyfht")
    def fht_inplace(x):
        N = len(x)
        i = N>>1
        while i:
            for k in range(0, N, 2*i):
                for j in range(k, i+k):
                    ij = i|j
                    temp = x[j]
                    x[j] += x[ij]
                    x[ij] = temp - x[ij]
            i = i >> 1


def sub_fht(n, m, seed=0, ordering=None):
    """
    Returns functions to compute the sub-sampled Walsh-Hadamard transform,
    i.e., operating with a wide rectangular matrix of random +/-1 entries.

    n: number of rows
    m: number of columns

    It is most efficient (but not required) for max(m+1,n+1) to be a power of 2.

    seed: determines choice of random matrix
    ordering: optional n-long array of row indices in [1, max(m,n)] to
              implement subsampling; generated by seed if not specified,
              but may be given to speed up subsequent runs on the same matrix.

    Returns (Ax, Ay, ordering):
        Ax(x): computes A.x (of length n), with x having length m
        Ay(y): computes A'.y (of length m), with y having length n
        ordering: the ordering in use, which may have been generated from seed
    """
    assert n > 0, "n must be positive"
    assert m > 0, "m must be positive"
    w = 2**int(np.ceil(np.log2(max(m+1, n+1))))

    if ordering is not None:
        assert ordering.shape == (n,)
    else:
        # in my code this function is never called without an ordering. So it doesn't really matter what this seed is. 
        rng = np.random.RandomState(seed)
        idxs = np.arange(1, w, dtype=np.uint32)
        rng.shuffle(idxs)
        ordering = idxs[:n]

    def Ax(x):
        assert x.size == m, "x must be m long"
        y = np.zeros(w)
        y[w-m:] = x.reshape(m)
        fht_inplace(y)
        return y[ordering]

    def Ay(y):
        assert y.size == n, "input must be n long"
        x = np.zeros(w)
        x.flat[ordering] = y
        fht_inplace(x)
        return x[w-m:]

    return Ax, Ay, ordering

def block_sub_fht(n, m, l, seed=0, ordering=None):
    #print("seed in block_sub_fht is: ", seed)
    """
    As `sub_fht`, but computes in `l` blocks of size `n` by `m`, potentially
    offering substantial speed improvements.

    n: number of rows
    m: number of columns per block
    l: number of blocks

    It is most efficient (though not required) when max(m+1,n+1) is a power of 2.

    seed: determines choice of random matrix
    ordering: optional (l, n) shaped array of row indices in [1, max(m, n)] to
              implement subsampling; generated by seed if not specified, but
              may be given to speed up subsequent runs on the same matrix.

    Returns (Ax, Ay, ordering):
        Ax(x): computes A.x (of length n), with x having length l*m
        Ay(y): computes A'.y (of length l*m), with y having length n
        ordering: the ordering in use, which may have been generated from seed
    """
    assert n > 0, "n must be positive"
    assert m > 0, "m must be positive"
    assert l > 0, "l must be positive"

    if ordering is not None:
        assert ordering.shape == (l, n)
    else:
        w = 2**int(np.ceil(np.log2(max(m+1, n+1))))
        # could just remove the random seed from here. 
        rng = np.random.RandomState(seed)
        ordering = np.empty((l, n), dtype=np.uint32)
        idxs = np.arange(1, w, dtype=np.uint32)
        for ll in range(l):
            rng.shuffle(idxs)
            ordering[ll] = idxs[:n]
        

    def Ax(x):
        assert x.size == l*m
        out = np.zeros(n)
        for ll in range(l):
            ax, ay, _ = sub_fht(n, m, ordering=ordering[ll])
            out += ax(x[ll*m:(ll+1)*m])
        return out

    def Ay(y):
        assert y.size == n
        out = np.empty(l*m)
        for ll in range(l):
            ax, ay, _ = sub_fht(n, m, ordering=ordering[ll])
            out[ll*m:(ll+1)*m] = ay(y)
        return out

    return Ax, Ay, ordering

#SPARC dictionary
#Returns 2 functions Ab and Az which compute Aβ and A^Tz respectively.
def sparc_transforms(L, M, n, seed=0):
    #print("seed in sparc transforms is: ", seed)
    Ax, Ay, ordering = block_sub_fht(n, M, L, ordering=None, seed=seed)
    def Ab(b):
        return Ax(b).reshape(-1, 1) / np.sqrt(n)
    def Az(z):
        return Ay(z).reshape(-1, 1) / np.sqrt(n)
    return Ab, Az, ordering
# return the ordering
# have second function which does the same but uses the previous ordering. And on fewer sections. 
# Check it on some betas. Use same beta and check it's the same result. Set last sections of beta to zero 
# For the second function just use the first L unprotected rows of the ordering. 


def sparc_transforms_shorter(L, M, n, ordering):
    """
    Returns 2 functions Ab and Az which compute Aβ and A^Tz respectively.
    Uses ordering to generate a smaller design matrix which is a subset of the original design matrix
    This allows the second round of amp decoding to be run on just the unprotected sections
    L: the number of unprotected sections (these must be at the start of the beta vector)
    """
    #print("calling block_sub_fht with an ordering, so doesn't matter what seed is ")
    # just use the first L unprotected rows of ordering
    Ax, Ay, _ = block_sub_fht(n, M, L, ordering=ordering[:L,:])
    def Ab(b):
        return Ax(b).reshape(-1, 1) / np.sqrt(n)
    def Az(z):
        return Ay(z).reshape(-1, 1) / np.sqrt(n)
    return Ab, Az


# Not using this currently, wasn't shown to make a difference so just using uniform power allocation. 
def pa_parameterised(L, C, P, a, f):
    """
    Power allocation
    Just going to include parameterised power allocation (but note there are two other options.)
    This results in an exponential pa with parameterised steepnees which is
    flattened after some point. 
    We generate a power allocation parameterised by a and f which set 
    the steepness and the flattening point respectively.
    exponential pa up to fL and then flat. The power is scaled to sum to P
    C: channel capacity
    """
    pa = 2**(-2 * a * C * np.arange(L) / L)
    pa[int(f*L):] = pa[int(f*L)]
    pa /= pa.sum() / P
    return pa

# the amp algorithm
def amp(y, σ_n, Pl, L, M, T, Ab, Az, β=np.array([None])):
    P = np.sum(Pl)
    n = y.size
    if β.all()==None:
        β = np.zeros((L*M, 1))
        z = y
        last_τ = 0
    else:
        β = β.reshape(L*M,1)
        z = y - Ab(β)
        # just to ensure the if statement isn't true on the first loop
        last_τ = 0
    
    for t in range(T):
        τ = np.sqrt(np.sum(z**2)/n)
        if τ == last_τ:
        # can change params for isclose if need be. If stopping too early.
        #if np.isclose(τ, last_τ):
            #print("Stopping after t runs:", t)
            return β
        last_τ = τ

        
        # added astype to increase precision to avoid divide by zero in LLR
        s = (β + Az(z))#.astype(np.longdouble)
        rt_n_Pl = np.sqrt(n*Pl).repeat(M).reshape(-1, 1)
        u = s * rt_n_Pl / τ**2
        max_u = u.max()
        exps = np.exp(u - max_u)
        sums = exps.reshape(L, M).sum(axis=1).repeat(M).reshape(-1, 1)
        β = (rt_n_Pl * exps / sums).reshape(-1, 1)
        z = y - Ab(β) + (z/τ**2) * (P - np.sum(β**2) / n)
    
    return β

#########End of code from Adam###########


class SPARCParams:
    """
    SPARC parameters
    l: number of sections in the overall SPARC
    m: number of columns per section in the SPARC
    p: channel signal power 
    sigma: is the channel noise standard deviation
    r: user data rate
    t: maximum number of AMP iterations permitted
    """
    def __init__(self, L, M, sigma, p, r, t, a=None, f=None, C=None):
        self.L = L
        self.M = M
        self.sigma = sigma
        self.p = p
        self.r = r
        self.t = t
        self.a = a
        self.f = f
        self.C = C


# LDPC Params needed to set up one of Jossy's codes
class LDPCParams:
    def __init__(self, standard, r_ldpc, z, ptype='A'):
        self.standard = standard
        self.r_ldpc = r_ldpc
        self.z = z
        self.ptype = ptype

def sp2bp(β, L, M: "must be power of 2"):
    """
    function to convert sectionwise probabilities to bitwise probabilities
    note: my sp2bp takes the entire beta vector, whereas adams code works only on 1 section. Change this if neccessary
    β: the normalised section posterior probabilities (each section should sum to 1)
    L: no. sections
    M: size of each section
    note L may not equal the total number of sections for the sparc code. 
    normally just interested in the number of sections covered by the LDPC
    """
    #initialise numpy array of zeros for all of the bitwise posteriors
    p = np.zeros(int(np.log2(M)*L))
    #loop through the L sections.
    for a in range(L):
        β_l = β[a*M:((a+1)*M)]
        for logi in range(int(np.log2(M))):
            b = int((a+1)*np.log2(M) - logi - 1)
            i = 2**logi
            k = i
            while k<M:
                #note shift in range of j due to βl being an array which is indexed from 0 not 1. 
                for j in range(k, k+i):
                    p[b] = p[b] + β_l[j]
                k = k + 2*i
    return p

def bp2sp(v, L, M: "must be power of 2"):
    """
    function to convert bitwise probabilities to sectionwise probabilities
    It is not possible to do this exactly so this is an approximation assuming independence between bitwise probabilities
    v: bitwise probabilities
    L: no. sections in v
    M: no. bits in v
    """
    logm = int(np.log2(M))
    # sectionwise probabilities
    sp = np.zeros(L*M)

    for l in range(L):
        # bitwise probabilities for section l
        bp = v[l*logm:(l+1)*logm]

        for m in range(M):
            # The bit values that correspond to the mth entry of this section being non-zero
            bits = bitarray(bin(m)[2:].zfill(logm))
            # bp and (1-bp) to the power of bits ensures that only the probability of each bit 
            # being the value we would need it to be for this column to be non-zero is included
            # in the probability. 
            # So we are multiplying the probability of each bit being how we need it to make the
            # current column non-zero.
            #print("bp", bp)
            #print("bits", bits)
            a = bp**bits
            bits.invert()
            b = (1-bp)**(bits)
            sp[l*M+m] = np.prod(a * b)
        sp[l*M:(l+1)*M] = sp[l*M:(l+1)*M]/sum(sp[l*M:(l+1)*M])
    return sp


def bits2indices(bits, m: "m must be a power of 2"):
    """
    assumes that the left most bit (lmb) is the most significant bit (msb)
    note that an input of all zero bits for one section will give an output of 0
    So the first column is being indexed by 0
    """
    logm = int(math.log(m,2))
    # assert that the bits can be split exactly into sections of size logm
    assert len(bits)%logm==0
    # find number of sections required for the bits
    sections = int(len(bits)/logm)
    # set up array to store the indices
    indices = []

    for i in range(sections):
        # get the bits corresponding to the current section
        section_bits = bits[i*logm:(i+1)*logm]
        # variable to store the value of the bits in indice
        digit = 0
        for j in range(len(section_bits)):
            if section_bits[j]:
                digit += 2**(len(section_bits)-j-1) 
        indices.append(digit)

    return indices

def ber_from_LLRs(M, LLR, input_indices, total_bits):
    """
    calculate the bit error rate for a set of LLRs compared to the original input indices
    M: columns per section of beta vector
    LLR: log likelihood ratios with possible errors
    input_indices: the indices we are trying to get with the LLRs
    total_bits: total numbers of bits the indices represent
    """
    # get hard decision on bits from the LLRs
    v_output = (LLR<0.0)
    # convert the bits to indices
    v_indices = bits2indices(v_output, M)
    # compute BER 
    return sum(bin(a^b).count('1') for (a, b) in zip(input_indices, v_indices))/total_bits

# note removed r_pa from function as never use it
def amp_ldpc_sim(sparcparams: SPARCParams, ldpcparams: LDPCParams = None, a=None, f=None, C=None):
    """
    Simulates original hard information exchange 
    Apply partial section coverage to a SPARC code. 
    Run AMP decoding on all, then LDPC decoding on the protected sections
    The protected sections are then hard decoded based on the output from the LDPC decoder and their 
    contribution is removed from the channel output. 
    A final round of AMP decoding is then run on the reduced channel output for the unprotected sections. 
    Can also be used for plain SPARC if ldpcparams=None
    a, f, and C control the parameterised power allocation 
    """
    #Get the SPARC parameters from the struct
    L = sparcparams.L
    M = sparcparams.M
    P = sparcparams.p
    sigma = sparcparams.sigma
    r_sparc = sparcparams.r
    #r_pa_sparc = sparcparams.r_pa
    T = sparcparams.t
    a = sparcparams.a
    f = sparcparams.f
    C = sparcparams.C

    # calculate some additional parameters 
    # Compute the SNR, capacity, and n, from the input parameters
    snr = P / sigma**2
    #C = 0.5 * np.log2(1 + snr)
    # n is number of rows of the design matrix 
    n = int(L*np.log2(M) / r_sparc)
    # compute additional parameters
    logm = np.log2(M)
    total_bits = int(logm*L)
    
    # Generate the power allocation
    if a==None: 
        # uniform power allocation across sections
        Pl = P/L * np.ones(L)
    else:
        Pl = pa_parameterised(L, C, P, a, f)

    if ldpcparams == None:
        # set params so that no ldpc encoding or decoding is used
        nl = 0
        kl = 0
        ldpc_bits = []
    else: 
        standard = ldpcparams.standard
        r_ldpc = ldpcparams.r_ldpc
        z = ldpcparams.z
        ptype = ldpcparams.ptype
        # initialise the ldpc code
        ldpc_code = ldpc.code(standard, r_ldpc, z, ptype)
        nl = ldpc_code.N
        kl = ldpc_code.K

        assert nl<=(L*logm)
        # we want the ldpc to cover a complete number of sections and not just part of the final section it covers. 
        assert nl%logm == 0
        # Generate random message in 2 stages
        # First generate ldpc bits
        protected_bits = np.random.randint(0, 2, kl).tolist()
        # Encode the protected_bits using the ldpc code 
        ldpc_bits = ldpc_code.encode(protected_bits).tolist() 

    # Generate the remaining bits required
    unprotected_bits = np.random.randint(0, 2, int(total_bits-nl)).tolist()
    # concatenate the ldpc and unprotected bits      
    sparc_bits = unprotected_bits+ldpc_bits
    assert len(sparc_bits)==total_bits
    # convert the bits to indices for encoding
    sparc_indices = bits2indices(sparc_bits, M)
    assert len(sparc_indices)==L
       
    # Generate the SPARC transform functions A.beta and A'.z
    Ab, Az, ordering = sparc_transforms(L, M, n)
    
    # Generate our transmitted signal x
    β_0 = np.zeros((L*M, 1))
    for l in range(L):
        β_0[l*M + sparc_indices[l]] = np.sqrt(n * Pl[l])
    x = Ab(β_0)
    
    # check that the power has been allocated uniformly. This should be approx equal to one when divided by the snr. 
    #print("Average power/snr is ", np.mean(x**2)/snr)

    # Generate random channel noise and then received signal y
    z = np.random.randn(n, 1) * sigma
    y = (x + z).reshape(-1, 1)
        
    # Run AMP decoding
    β = amp(y, sigma, Pl, L, M, T, Ab, Az).reshape(-1)
    
     # Convert decoded beta back to a message
    rx_message = []
    for l in range(L):
        idx = np.argmax(β[l*M:(l+1)*M])
        rx_message.append(idx)
    
    # Compute fraction of sections decoded correctly with just AMP decoding
    #correct_amp = np.sum(np.array(rx_message) == np.array(sparc_indices)) / L
    #print("Fraction of sections decoded correctly with amp: ", correct_amp)
    
    # Compute BER for just amp decoding
    ber_amp = sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, rx_message))/total_bits

    if ldpcparams==None:
        ber_ldpc = None
        ber_ldpc_amp = None
    else:
        # Get the sectionwise posterior probabilities by dividing β by the power in each section. 
        # This converts each section in β to a valid probability distribution.
        sectionwise_posterior = β/np.sqrt(n*np.repeat(Pl, M))

        ldpc_sections = int(nl/logm)
        # convert sectionwise to bitwise posterior probabilities
        # only need the sections corresponding to ldpc code 
        bitwise_posterior = sp2bp(sectionwise_posterior[(L-ldpc_sections)*M:], ldpc_sections, M) 
        # computer the log likelihood ratio for decoding
        LLR = np.log(1-bitwise_posterior)- np.log(bitwise_posterior)
        # set -inf and +inf to real numbers with v large magnitude
        LLR = np.nan_to_num(LLR)

        (app, it) = ldpc_code.decode(LLR)

        #test_output = (bitwise_posterior>0.5)
        #print("Incorrect bits before ldpc decoding: ", np.sum(test_output!=ldpc_bits))

        v_output = (app<0.0)
        #print("Incorrect bits in ldpc decoding: ", np.sum(v_output!=ldpc_bits))

        # convert the bits to indices
        beta_output_ldpc = bits2indices(v_output, M)

        beta_output = rx_message
        beta_output[L-ldpc_sections:] = beta_output_ldpc

        # compute fraction of sections decoded correctly with AMP followed by LDPC decoding
        #correct_ldpc = np.sum(np.array(beta_output) == np.array(sparc_indices))/L
        #print("Fractions of sections decoded correctly with amp and ldpc: ", correct_ldpc)

        # compute BER for amp and ldpc
        ber_ldpc = sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, beta_output))/total_bits

        # only perform final round of AMP decoding if not all sections of beta are covered with LDPC
        if (L-ldpc_sections==0):
            ber_ldpc_amp = None
        else:    
            # generate beta_ldpc which has the first L unprotected sections set to zero
            # and the L ldpc sections will have exactly one non-zero entry per section. 
            beta_ldpc = np.zeros((L*M, 1))
            # using variable i because beta_output_L only has ldpc_sections number of entries. 
            i=0
            for l in range(L-ldpc_sections,L):
                beta_ldpc[(l)*M + beta_output_ldpc[i]] = np.sqrt(n * Pl[l])
                i+=1
            x_ldpc = Ab(beta_ldpc)

            # calculate a new value of the channel output which doesn't contain the contribution from 
            # the ldpc sections.
            y_new = y - x_ldpc

            # Rerun amp decoding on the new input y_new over the unprotected sections.
            L_unprotected = L-ldpc_sections
            Ab_new, Az_new = sparc_transforms_shorter(L_unprotected, M, n, ordering)

            beta_new = amp(y_new, sigma, Pl[:L_unprotected], L_unprotected, M, T, Ab_new, Az_new).reshape(-1)

            # Convert decoded beta_new back to a message
            # beta_new gives you the first L-ldpc_sections sections of the final received message. 
            rx_message_final = []
            for l in range(L_unprotected):
                idx = np.argmax(beta_new[l*M:(l+1)*M])
                rx_message_final.append(idx)

            beta_output[:L_unprotected] = rx_message_final
            # compute fraction of sections decoded correctly with AMP followed by SPARC decoding followed by amp.
            #correct_ldpc = np.sum(np.array(beta_output) == np.array(sparc_indices))/L
            #print("Fractions of sections decoded correctly with amp and ldpc and amp: ", correct_ldpc)

            # compute BER for amp and ldpc and amp again.
            ber_ldpc_amp = sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, beta_output))/total_bits      
    # overall rate of the code
    R = (L*logm - (nl-kl))/n
    #Compute Eb/N0
    #EbN0 = 1/(2*R) * (P/sigma**2)
   
    return ber_amp, ber_ldpc, ber_ldpc_amp, R

def soft_amp_ldpc_sim(sparcparams: SPARCParams, ldpcparams: LDPCParams, soft_iter, a=None, f=None, C=None):
    """
    Simulates soft information exchange. 
    Soft output from LDPC used to initialise beta_0 in the amp decoding. 
    In this function '_ldpc' means after ldpc decoding and '_amp' means after amp decoding. 
    soft_iter: the number of iterations of soft information exchange. 
    a, f, and C control the parameterised power allocation
    """
    # arrays to store the return values of the ber
    # bit error rate after each round of amp decoding
    ber_amp = []
    # ber after each round of ldpc decoding
    ber_ldpc = []

    #Get the SPARC parameters from the struct
    L = sparcparams.L
    M = sparcparams.M
    P = sparcparams.p
    sigma = sparcparams.sigma
    r_sparc = sparcparams.r
    #r_pa_sparc = sparcparams.r_pa
    T = sparcparams.t
    a = sparcparams.a
    f = sparcparams.f
    C = sparcparams.C

    # calculate some additional parameters 
    # Compute the SNR, capacity, and n, from the input parameters
    snr = P / sigma**2
    #C = 0.5 * np.log2(1 + snr)
    # n is number of rows of the design matrix 
    n = int(L*np.log2(M) / r_sparc)
    # compute additional parameters
    logm = np.log2(M)
    total_bits = int(logm*L)
    
    # Generate the power allocation
    if a==None:
        # uniform power allocation across sections
        Pl = P/L * np.ones(L)
    else:
        # parameterised power allocation
        Pl = pa_parameterised(L, C, P, a, f)

    standard = ldpcparams.standard
    r_ldpc = ldpcparams.r_ldpc
    z = ldpcparams.z
    ptype = ldpcparams.ptype
    # initialise the ldpc code
    ldpc_code = ldpc.code(standard, r_ldpc, z, ptype)
    nl = ldpc_code.N
    kl = ldpc_code.K

    assert nl<=(L*logm)
    # we want the ldpc to cover a complete number of sections and not just part of the final section it covers. 
    assert nl%logm == 0

    # Generate random message in 2 stages
    # First generate ldpc bits
    protected_bits = np.random.randint(0, 2, kl).tolist()
    # Encode the protected_bits using the ldpc code 
    ldpc_bits = ldpc_code.encode(protected_bits).tolist() 
    # Generate the remaining bits required
    unprotected_bits = np.random.randint(0, 2, int(total_bits-nl)).tolist()
    # concatenate the ldpc and unprotected bits      
    sparc_bits = unprotected_bits+ldpc_bits

    assert len(sparc_bits)==total_bits

    # convert the bits to indices for encoding
    sparc_indices = bits2indices(sparc_bits, M)
    assert len(sparc_indices)==L
       
    # Generate the SPARC transform functions A.beta and A'.z
    Ab, Az, ordering = sparc_transforms(L, M, n)
    
    # Generate our transmitted signal x
    β_0 = np.zeros((L*M, 1))
    for l in range(L):
        β_0[l*M + sparc_indices[l]] = np.sqrt(n * Pl[l])
    x = Ab(β_0)
    
    # check that the power has been allocated uniformly. This should be approx equal to one when divided by the snr. 
    #print("Average power/snr is ", np.mean(x**2)/snr)

    # Generate random channel noise and then received signal y
    z = np.random.randn(n, 1) * sigma
    y = (x + z).reshape(-1, 1)
    
    # Run AMP decoding
    β = amp(y, sigma, Pl, L, M, T, Ab, Az).reshape(-1)
    
    # Convert decoded beta back to a message
    rx_message = []
    for l in range(L):
        idx = np.argmax(β[l*M:(l+1)*M])
        rx_message.append(idx)
    
    # Compute fraction of sections decoded correctly with just AMP decoding
    #correct_amp = np.sum(np.array(rx_message) == np.array(sparc_indices)) / L
    #print("Fraction of sections decoded correctly with amp: ", correct_amp)
    
    # Compute BER for just sparc
    ber_amp.append(sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, rx_message))/total_bits)

    beta_amp = β
    rx_message_amp = rx_message
    for i in range(soft_iter):
        # Get the sectionwise posterior probabilities by dividing β by the power in each section. 
        # This converts each section in β to a valid probability distribution.
        sectionwise_posterior = beta_amp/np.sqrt(n*np.repeat(Pl, M))

        #print('sectionwise posterior sum ', sum(sectionwise_posterior))

        # remember only want to perform ldpc decoding on sections covered by the ldpc code
        ldpc_sections = int(nl/logm)
        # convert sectionwise to bitwise posterior probabilities
        # only need the sections corresponding to ldpc code 
        bitwise_posterior = sp2bp(sectionwise_posterior[(L-ldpc_sections)*M:], ldpc_sections, M) 
        # compute the log likelihood ratio for decoding
        LLR = np.log(1-bitwise_posterior)- np.log(bitwise_posterior)
        # set -inf and +inf to real numbers with v large magnitude
        LLR = np.nan_to_num(LLR)
        (app, it) = ldpc_code.decode(LLR)
        # get hard decision on bits after ldpc
        v_output = (app<0.0)
        # convert the bits to indices
        v_indices = bits2indices(v_output, M)

        # replace old indices for ldpc sections with new ones
        ldpc_indices = rx_message_amp
        ldpc_indices[L-ldpc_sections:] = v_indices

        # compute BER after ldpc decoding
        ber_ldpc.append(sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, ldpc_indices))/total_bits)

        # now perform amp decoding using the soft output from the ldpc
        # get the bitwise a posterior probabilities from the a posterior log likelihood ratios in app
        bitwise_ldpc = 1/(1+np.exp(app))

        # check bitwise_ldpc is made up of valid probabilities
        assert(bitwise_ldpc.all()>=0 and bitwise_ldpc.all()<=1)

        # convert the bitwise posterior to sectionwise posterior
        sectionwise_ldpc = sectionwise_posterior
        sectionwise_ldpc[M*(L-ldpc_sections):] = bp2sp(bitwise_ldpc, ldpc_sections, M)
        #print("sum sectionwise_ldpc ", sum(sectionwise_ldpc))

        # multiply by the power allocation
        beta_ldpc = sectionwise_ldpc * np.sqrt(n*np.repeat(Pl, M))
        # Run AMP decoding initialising the amp decoder with beta_ldpc
        beta_amp = amp(y, sigma, Pl, L, M, T, Ab, Az, beta_ldpc).reshape(-1)

        # Convert decoded beta_amp back to a message 
        rx_message_amp = []
        for l in range(L):
            idx = np.argmax(beta_amp[l*M:(l+1)*M])
            rx_message_amp.append(idx)

        ber_amp.append(sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, rx_message_amp))/total_bits) 

    # overall rate of the code
    R = (L*logm - (nl-kl))/n
    #Compute Eb/N0
    #EbN0 = 1/(2*R) * (P/sigma**2)
    return ber_amp, ber_ldpc, R


def hardinitbeta_amp_ldpc_sim(sparcparams: SPARCParams, ldpcparams: LDPCParams):
    """
    Simulates hard information exchange with multiple iterations. 
    Output from LDPC decoder is hard decoded to give a beta vector. This ahrd decoded beta vector will then be used 
    as the initialisation of beta_0 in the amp decoder. 
    In this function '_ldpc' means after ldpc decoding and '_amp' means after amp decoding
    Only offers uniform power allocation
    """
    # arrays to store the return values of the ber
    # bit error rate after each round of amp decoding
    ber_amp = []
    # ber after each round of ldpc decoding
    ber_ldpc = []

    #Get the SPARC parameters from the struct
    L = sparcparams.L
    M = sparcparams.M
    P = sparcparams.p
    sigma = sparcparams.sigma
    r_sparc = sparcparams.r
    #r_pa_sparc = sparcparams.r_pa
    T = sparcparams.t

    # calculate some additional parameters 
    # Compute the SNR, capacity, and n, from the input parameters
    snr = P / sigma**2
    #C = 0.5 * np.log2(1 + snr)
    # n is number of rows of the design matrix 
    n = int(L*np.log2(M) / r_sparc)
    # compute additional parameters
    logm = np.log2(M)
    total_bits = int(logm*L)
    
    # Generate the power allocation
    # uniform power allocation across sections
    Pl = P/L * np.ones(L)

    standard = ldpcparams.standard
    r_ldpc = ldpcparams.r_ldpc
    z = ldpcparams.z
    ptype = ldpcparams.ptype
    # initialise the ldpc code
    ldpc_code = ldpc.code(standard, r_ldpc, z, ptype)
    nl = ldpc_code.N
    kl = ldpc_code.K

    assert nl<=(L*logm)
    # we want the ldpc to cover a complete number of sections and not just part of the final section it covers. 
    assert nl%logm == 0

    # Generate random message in 2 stages
    # First generate ldpc bits
    protected_bits = np.random.randint(0, 2, kl).tolist()
    # Encode the protected_bits using the ldpc code 
    ldpc_bits = ldpc_code.encode(protected_bits).tolist() 
    # Generate the remaining bits required
    unprotected_bits = np.random.randint(0, 2, int(total_bits-nl)).tolist()
    # concatenate the ldpc and unprotected bits      
    sparc_bits = unprotected_bits+ldpc_bits
    assert len(sparc_bits)==total_bits
    # convert the bits to indices for encoding
    sparc_indices = bits2indices(sparc_bits, M)
    assert len(sparc_indices)==L
       
    # Generate the SPARC transform functions A.beta and A'.z
    Ab, Az, ordering = sparc_transforms(L, M, n)
    
    # Generate our transmitted signal x
    β_0 = np.zeros((L*M, 1))
    for l in range(L):
        β_0[l*M + sparc_indices[l]] = np.sqrt(n * Pl[l])
    x = Ab(β_0)
    
    # check that the power has been allocated uniformly. This should be approx equal to one when divided by the snr. 
    #print("Average power/snr is ", np.mean(x**2)/snr)

    # Generate random channel noise and then received signal y
    z = np.random.randn(n, 1) * sigma
    y = (x + z).reshape(-1, 1)
    
    # Run AMP decoding
    β = amp(y, sigma, Pl, L, M, T, Ab, Az).reshape(-1)
    
    # Convert decoded beta back to a message
    rx_message = []
    for l in range(L):
        idx = np.argmax(β[l*M:(l+1)*M])
        rx_message.append(idx)
    
    # Compute fraction of sections decoded correctly with just AMP decoding
    #correct_amp = np.sum(np.array(rx_message) == np.array(sparc_indices)) / L
    #print("Fraction of sections decoded correctly with amp: ", correct_amp)
    
    # Compute BER for just sparc
    ber_amp.append(sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, rx_message))/total_bits)

    beta_amp = β
    rx_message_amp = rx_message
    # Get the sectionwise posterior probabilities by dividing β by the power in each section. 
    # This converts each section in β to a valid probability distribution.
    sectionwise_posterior = beta_amp/np.sqrt(n*np.repeat(Pl, M))
    #print('sectionwise posterior sum ', sum(sectionwise_posterior))
    # remember only want to perform ldpc decoding on sections covered by the ldpc code
    ldpc_sections = int(nl/logm)
    # convert sectionwise to bitwise posterior probabilities
    # only need the sections corresponding to the ldpc code 
    bitwise_posterior = sp2bp(sectionwise_posterior[(L-ldpc_sections)*M:], ldpc_sections, M) 
    # compute the log likelihood ratio for decoding
    LLR = np.log(1-bitwise_posterior)- np.log(bitwise_posterior)
    # set -inf and +inf to real numbers with v large magnitude
    LLR = np.nan_to_num(LLR)
    (app, it) = ldpc_code.decode(LLR)
    # get hard decision on bits after ldpc
    v_output = (app<0.0)
    # convert the bits to indices
    v_indices = bits2indices(v_output, M)

    # replace old indices for ldpc sections with new ones
    ldpc_indices = rx_message_amp
    ldpc_indices[L-ldpc_sections:] = v_indices

    # compute BER after ldpc decoding
    ber_ldpc.append(sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, ldpc_indices))/total_bits)

    # use the hard decoded indices to produce a hard decoded beta from the LDPC decoding
    beta_ldpc = np.zeros((L*M,1))
    for l in range(L):
        # this will set one position in each section to the appropriate power allocation.
        beta_ldpc[l*M + ldpc_indices[l]] = np.sqrt(n * Pl[l])

    # Run AMP decoding with the original channel output and this new intialisation for beta_0
    beta_amp = amp(y, sigma, Pl, L, M, T, Ab, Az, beta_ldpc).reshape(-1)

    # Convert decoded beta_amp back to a message 
    rx_message_amp = []
    for l in range(L):
        idx = np.argmax(beta_amp[l*M:(l+1)*M])
        rx_message_amp.append(idx)

    ber_amp.append(sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, rx_message_amp))/total_bits) 

    # overall rate of the code
    R = (L*logm - (nl-kl))/n
    #Compute Eb/N0
    #EbN0 = 1/(2*R) * (P/sigma**2)
    return ber_amp, ber_ldpc, R

def soft_amp_ldpc_hardinit(sparcparams: SPARCParams, ldpcparams: LDPCParams, soft_iter, threshold):
    """
    Simulates threshold initialised information exchange.
    If an ldpc section has a probability greater than the threshold it is hard decoded and it's contribution 
    removed from the channel output. AMP decoding is then run on the reduced channel output on sections that aren't hard decoded
    The following round of LDPC decoding uses the LLRs from before AMP decoding for the hard decoded section. 
    This allows information exchange to continue going back and forward between the AMP decoder and LDPC decoder. 
    Only offers uniform power allocation 
    soft_iter: the number of iterations of information exchange
    threshold: the threshold used for hard decoding the ldpc decoder output. 
    """
    # arrays to store the return values of the ber
    # bit error rate after each round of amp decoding
    ber_amp = []
    # ber after each round of ldpc decoding
    ber_ldpc = []

    #Get the SPARC parameters from the struct
    L = sparcparams.L
    M = sparcparams.M
    P = sparcparams.p
    sigma = sparcparams.sigma
    r_sparc = sparcparams.r
    #r_pa_sparc = sparcparams.r_pa
    T = sparcparams.t

    # calculate some additional parameters 
    # Compute the SNR, capacity, and n, from the input parameters
    snr = P / sigma**2
    #C = 0.5 * np.log2(1 + snr)
    # n is number of rows of the design matrix 
    n = int(L*np.log2(M) / r_sparc)
    # compute additional parameters
    logm = int(np.log2(M))
    total_bits = int(logm*L)
    
    # Generate the power allocation
    # uniform power allocation across sections
    Pl = P/L * np.ones(L)
    
    standard = ldpcparams.standard
    r_ldpc = ldpcparams.r_ldpc
    z = ldpcparams.z
    ptype = ldpcparams.ptype
    # initialise the ldpc code\
    ldpc_code = ldpc.code(standard, r_ldpc, z, ptype)
    nl = ldpc_code.N
    kl = ldpc_code.K

    assert nl<=(total_bits)
    # we want the ldpc to cover a complete number of sections and not just part of the final section it covers. 
    assert nl%logm == 0
    #(L, M, sigma, P, R, T, R_PA, R_LDPC):

    if standard=="802.11n" or standard=="802.16":
        # Generate random message in 2 stages
        # First generate ldpc bits
        protected_bits = np.random.randint(0, 2, kl).tolist()
        assert len(protected_bits) == kl
        # Encode the protected_bits using the ldpc code 
        ldpc_bits = ldpc_code.encode(protected_bits).tolist() 
        # Generate the remaining bits required
        unprotected_bits = np.random.randint(0, 2, int(total_bits-nl)).tolist()
        # concatenate the ldpc and unprotected bits      
        sparc_bits = unprotected_bits+ldpc_bits
        assert len(sparc_bits)==total_bits
        # seed for the design matrix A. We want this to be the same everytime
        seed=0
    else:
        # use the all zero codeword for the sparc bits to avoid the need for ldpc encoding.
        sparc_bits = np.zeros(total_bits)
        # don't want to set a random seed. This should mean the function seeds off the clock time or similar.
        # This introduces randomness (which this time is not provided by the initialisation of the sparc_bits)
        seed=None
    # convert the bits to indices for encoding
    sparc_indices = bits2indices(sparc_bits, M)
    assert len(sparc_indices)==L
       
    # Generate the SPARC transform functions A.beta and A'.z
    Ab, Az, ordering = sparc_transforms(L, M, n, seed)
    # Generate our transmitted signal X
    β_0 = np.zeros((L*M, 1))
    for l in range(L):
        β_0[l*M + sparc_indices[l]] = np.sqrt(n * Pl[l])
    x = Ab(β_0)
    # check that the power has been allocated uniformly. This should be approx equal to one when divided by the snr. 
    #print("Average power/snr is ", np.mean(x**2)/snr)
    # Generate random channel noise and then received signal y
    z = np.random.randn(n, 1) * sigma
    y = (x + z).reshape(-1, 1)
    
    # Run AMP decoding
    β = amp(y, sigma, Pl, L, M, T, Ab, Az).reshape(-1)
    # Convert decoded beta back to a message
    rx_message = []
    for l in range(L):
        idx = np.argmax(β[l*M:(l+1)*M])
        rx_message.append(idx)
    
    # Compute fraction of sections decoded correctly with just AMP decoding
    #correct_amp = np.sum(np.array(rx_message) == np.array(sparc_indices)) / L
    #print("Fraction of sections decoded correctly with amp: ", correct_amp)
    
    # Compute BER for just sparc
    ber_amp.append(sum(bin(a^b).count('1') for (a, b) in zip(sparc_indices, rx_message))/total_bits)

    beta_amp = β
    rx_message_amp = rx_message
    # Get the sectionwise posterior probabilities by dividing β by the power in each section. 
    # This converts each section in β to a valid probability distribution.
    sectionwise_posterior = beta_amp/np.sqrt(n*np.repeat(Pl, M))
    #print('sectionwise posterior sum ', sum(sectionwise_posterior))
    # remember only want to perform ldpc decoding on sections covered by the ldpc code
    ldpc_sections = int(nl/logm)
    # convert sectionwise to bitwise posterior probabilities for all sections
    bitwise_posterior = sp2bp(sectionwise_posterior, L, M) 
    # compute the log likelihood ratio for decoding
    LLR = np.log(1-bitwise_posterior)- np.log(bitwise_posterior)
    # set -inf and +inf to real numbers with v large magnitude
    LLR = np.nan_to_num(LLR)
    for i in range(soft_iter):
        # perform ldpc decoding on the LDPC sections
        (app, it) = ldpc_code.decode(LLR[(L-ldpc_sections)*logm:])
        # Replace the ldpc LLRS in the LLRs from AMP decoding
        LLR[(L-ldpc_sections)*logm:] = app
        # compute BER after ldpc decoding
        ber_ldpc.append(ber_from_LLRs(M, LLR, sparc_indices, total_bits))
        # want to finish the soft iterations on LDPC decoding. 
        if i==(soft_iter - 1):
            break
        
        # now perform amp decoding using the soft output from the ldpc
        # get the bitwise a posterior probabilities from the a posterior log likelihood ratios in app
        bitwise_ldpc = 1/(1+np.exp(app))
        # check bitwise_ldpc is made up of valid probabilities
        assert(bitwise_ldpc.all()>=0 and bitwise_ldpc.all()<=1)

        # convert the bitwise posterior to sectionwise posterior
        sectionwise_ldpc = sectionwise_posterior
        sectionwise_ldpc[M*(L-ldpc_sections):] = bp2sp(bitwise_ldpc, ldpc_sections, M)
        #print("sum sectionwise_ldpc ", sum(sectionwise_ldpc))

        # produce the hard initialisation for the amp decoding
        # Will then only perform amp decoding on sections which don't have an entry in their section with probability exceeding the threshold. 
        # All sections not covered by the ldpc code will always be involved in the amp decoding. 
        y_new, Ab_new, Az_new, amp_sections, L_amp_sections = ae.hard_initialisation(sectionwise_ldpc, L, M, n, ordering, y, Pl, Ab, threshold, ldpc_sections)
        # NOTE HARD CODED IN THE RATE IN PRINT STATEMENTS FOR SPEED
        #print("EbN0 in dB is (with hard coded rate): ", 20*np.log10(snr/(2*0.5)))
        #print("iteration is: ", i)
        #print("L_amp_sections: ", L_amp_sections)

        # keep the LLRs corresponding to the sections hard decoded before performing amp
        # if there are no amp sections we don't change LLR
        if L_amp_sections>0:
            # perform amp decoding on the amp sections
            beta_T = amp(y_new, sigma, Pl[amp_sections], L_amp_sections, M, T, Ab_new, Az_new).reshape(-1)

            # convert sectionwise posterior probabilities to bitwise posterior probabilities
            sectionwise = beta_T/np.sqrt(n*np.repeat(Pl[amp_sections],M))

            #print("sum of sectionwise is: ", np.sum(sectionwise))
            #print("sum of sectionwise should be: ", L_amp_sections)
            bitwise = sp2bp(sectionwise, L_amp_sections, M)
            # convert bitwise post. to LLRs 
            LLR_amp_sections = np.log(1-bitwise)-np.log(bitwise)
            # set -inf and +inf to real numbers with v large magnitude
            LLR_amp_sections=np.nan_to_num(LLR_amp_sections)
        
            # convert amp_sections to a list of all the positions in these sections
            amp_positions = np.tile(np.linspace(0,logm-1,logm,dtype=np.dtype(np.int16)),L_amp_sections)
            amp_positions = amp_positions + logm*np.repeat(amp_sections, logm)
            # replace the LLRs for the amp_sections with the new LLRs
            LLR[amp_positions] = LLR_amp_sections
            # use sectionwise_posterior to store the sectionwise posteriors of the sections which aren't protected by the ldpc
            # Note this will always be the first (L-ldpc_sections)*M columns of sectionwise
            # actually this is redundant as we just replace these with each loop anyway. 
            sectionwise_posterior[:(L-ldpc_sections)*M] = sectionwise[:(L-ldpc_sections)*M]
        # calculate the new BER after amp decoding for this new set of LLRs
        ber_amp.append(ber_from_LLRs(M, LLR, sparc_indices, total_bits))
    # overall rate of the code
    R = (L*logm - (nl-kl))/n
    #Compute Eb/N0
    #EbN0 = 1/(2*R) * (P/sigma**2)
   
    return ber_amp, ber_ldpc, R


def awgn(x, sigma):
    # Code for simulating the AWGN channel   
    noise = sigma*np.random.randn(len(x))
    return x + noise

def ch2llr(ch, sigma):
    # convert channel output to LLR
    sigma2 = sigma**2
    return 2.0/sigma2*ch

def bpsk(x):
    # convert bits to bpsk (i.e. 0->1, 1->-1)
    return 1.0 - 2.0*x  


def sim_ldpc(ldpcparams: LDPCParams, sigma, MIN_ERRORS = 100, MAX_BLOCKS = 400000):
    """
    Simulation of using LDPC code with BPSK modulation on the AWGN channel. 
    sigma: standard deviation of the channel noise. 
    """
    
    rate = ldpcparams.r_ldpc
    standard = ldpcparams.standard
    z = ldpcparams.z
    ptype = ldpcparams.ptype

    if rate == "1/2":
        R = .5
    elif rate == "2/3":
        R = 0.6667
    elif rate == "3/4":
        R = 0.75
    elif rate == "5/6":
        R = 0.83333
    elif rate=="0.45":
        R=0.45
    else:
        raise NameError("Rate unsupported")

    mycode = ldpc.code(standard, rate, z, ptype)
    K = mycode.K
    N = mycode.N


    nbiterrors = 0
    nblockerrors = 0
    nblocks = 0
    nit_total = 0
    while nblockerrors < MIN_ERRORS:
        if standard=="802.11n" or standard=="802.16":
            u = np.random.randint(0,2,K)
            x = mycode.encode(u)
        else:
            x=np.zeros(N)
        xm = bpsk(x)

        y = awgn(xm, sigma)

        yl = ch2llr(y, sigma)
        (app,nit) = mycode.decode(yl, 'sumprod2')
        xh = (app < 0.0)

        biterrors_thisblock = np.sum(x != xh)
        nbiterrors += biterrors_thisblock
        if biterrors_thisblock:
            nblockerrors += 1
        nblocks += 1
        nit_total += nit

        # errors per block divided by the total number of information bits.
        ber = nbiterrors/(nblocks*N)

        if nblocks >= MAX_BLOCKS:
            break

    return ber

def waterfall(sparcparams: SPARCParams, ldpcparams: LDPCParams, csv_filename: str, png_filename: str, init='soft', pa_param=False, datapoints=10, MIN_ERRORS=100, MAX_BLOCKS=500, bpsk=True, sections=512):
    """
    Code to plot a waterfall curve for an LDPC code with bpsk modulation and an LDPC code with sparc modulation. 
    init: determines which form of initialisation is used. Options are soft, hard, originalHard. 
    Note that threshold initialised information exchange is not an option here. 
    Note this fn is HARD CODED to work with rate 5/6 ldpc code. 
    sections: determines the section coverage. Should only be used with originalHard. Set to 512 as default as 
    it is assumed that L=512.
    Number of ldpc sections must be divisible by 8 to ensure nl is divisible by 24. 
    pa_param: determines which power allocation is used. pa_param=False means uniform power allocation. 
    """
    # Sparc parameters
    L = sparcparams.L
    M = sparcparams.M
    logm = np.log2(M)
    p=sparcparams.p
    r_sparc = sparcparams.r
    T = sparcparams.t
    a = sparcparams.a
    f = sparcparams.f
    C = sparcparams.C

    # LDPC parameters
    standard = ldpcparams.standard
    r_ldpc = ldpcparams.r_ldpc
    # number of ldpc sections, must be divisible by 8 to ensure nl is divisible by 24.
    # number of ldpc bits, must be divisible by 24
    nl = logm * sections
    z = int(nl/24)
    ldpcparams = LDPCParams(standard, r_ldpc, z)    

    n = L*logm/r_sparc
    # HARD CODED RATE 5/6
    print("Note that rate is hard coded as 5/6")
    R = (L*logm-nl*(1-5/6))/n
    # note that R should be 5/6 unless full section coverage is not used
    print('Overall rate is: ', R)

    # snr at capacity
    snrc = (2**(2*R) - 1)
    EbN0c = 1/(2*R) * snrc
    EbN0c_dB = 20*np.log10(EbN0c)

    EbN0_dB = np.linspace(3, 10, datapoints)
    # BER after first round of AMP decoding
    BER_amp_1 = np.zeros(datapoints)
    # BER after second round of AMP decoding
    BER_amp_2 = np.zeros(datapoints)
    # BER after first round of ldpc
    BER_ldpc = np.zeros(datapoints)
    # BER after 2nd round ldpc
    BER_ldpc_2 = np.zeros(datapoints)
    BER_bpsk = np.zeros(datapoints)
    # BER of plain SPARC
    BER_plain = np.zeros(datapoints)

    i=0
    for ebno_db in EbN0_dB:
        ebno = 10**(ebno_db/20)
        if bpsk==True:
            # sigma for the bpsk results
            # to calculate this we use use the fact that E_b=energy per bit=1 in this case
            # N_0/2 = sigma**2 where sigma**2 is the variance of the noise. 
            E_b = 1
            N_0 = E_b/ebno
            sigma_bpsk = np.sqrt(N_0/2)
            # get the BER for the ldpc and bpsk modulation. The repeats are built into the code
            BER_bpsk[i] = sim_ldpc(ldpcparams, sigma_bpsk, MIN_ERRORS, MAX_BLOCKS)

        # sigma for the sparc based methods
        #snrdB = 20*np.log10(P/sigma**2)
        #EbN0 = 1/(2*R) * (p/SIGMA**2)
        #EbN0_dB = 20*log10(EbN0)
        snr = ebno/(1/(2*R))
        sigma = np.sqrt(p/snr)
        C = 0.5*np.log2(1+p/(sigma**2))
        if pa_param==True and a==None:
            a=r_sparc/C
            f=r_sparc/C
        # Note that I'm using the same power allocation for both SPARCs, but it is being tailored to the sparc with outer code. 
        sparcparams = SPARCParams(L, M, sigma, p, r_sparc, T, a, f, C)
        # plain sparc with same overall rate for comparison
        sparcparams_plain = SPARCParams(L, M, sigma, p, R, T, a, f, C)
        # cumulative BER for amp and ldpc
        ber_cum_amp1 = 0
        ber_cum_amp2 = 0
        ber_cum_ldpc = 0
        ber_cum_ldpc2 = 0
        ber_cum_plain = 0
        nblockerrors_ldpc = 0
        nblocks = 0
        while nblockerrors_ldpc < MIN_ERRORS:
            if init=='soft':
                (ber_thisblock_amp, ber_thisblock_ldpc, _) = soft_amp_ldpc_sim(sparcparams, ldpcparams, 2)
            elif init=='hard':
                (ber_thisblock_amp, ber_thisblock_ldpc, _) = hardinitbeta_amp_ldpc_sim(sparcparams, ldpcparams)
                # append a zero so that the calculations for the soft init don't fail.
                ber_thisblock_ldpc.append(0)
            elif init=='originalHard':
                ber_thisblock_amp = np.zeros(2)
                # have zero in position 1 so calculations for the soft init don't fail
                ber_thisblock_ldpc = np.zeros(2)
                (ber_thisblock_amp[0], ber_thisblock_ldpc[0], ber_thisblock_amp[1], _) = amp_ldpc_sim(sparcparams, ldpcparams)      
            else:
                print(init, " is not a valid initialisation. Please change to \'soft\' or \'hard\'")
            (ber_thisblock_plain, _, _, _) = amp_ldpc_sim(sparcparams_plain)
            #print("BER amp:", ber_thisblock_amp)
            #print("BER ldpc: ", ber_thisblock_ldpc)

            ber_cum_amp1 += ber_thisblock_amp[0]
            ber_cum_amp2 += ber_thisblock_amp[1]
            ber_cum_ldpc += ber_thisblock_ldpc[0]
            ber_cum_ldpc2 += ber_thisblock_ldpc[1]
            ber_cum_plain += ber_thisblock_plain
            if ber_thisblock_plain:
                nblockerrors_ldpc += 1
            nblocks += 1

            if nblocks >= MAX_BLOCKS:
                break

        BER_amp_1[i] = ber_cum_amp1/nblocks
        BER_amp_2[i] = ber_cum_amp2/nblocks
        BER_ldpc[i] = ber_cum_ldpc/nblocks
        BER_ldpc_2[i] = ber_cum_ldpc2/nblocks
        BER_plain[i] = ber_cum_plain/nblocks


        i+=1
    
    # open file you want to write CSV output to. 'a' means its in append mode. Switching this to 'w' will make it overwrite the file.
    myFile = open(csv_filename, 'a')
    with myFile:
        # if bpsk False, will just have all zeros for BER_bpsk
        myFields = ['EbN0_dB', 'BER_amp_1', 'BER_ldpc', 'BER_amp_2', 'BER_ldpc_2', 'BER_plain', 'BER_bpsk']
        writer = csv.DictWriter(myFile, fieldnames=myFields)
        writer.writeheader()
        for k in range(datapoints):
            writer.writerow({'EbN0_dB' : EbN0_dB[k], 'BER_amp_1': BER_amp_1[k],'BER_ldpc' : BER_ldpc[k], 'BER_amp_2': BER_amp_2[k], 'BER_ldpc_2': BER_ldpc_2[k], 'BER_plain': BER_plain[k], 'BER_bpsk': BER_bpsk[k]})

    print("BER_ldpc is: ", BER_ldpc)
    fig, ax = plt.subplots()
    ax.set_yscale('log', basey=10)
    ax.plot(EbN0_dB, BER_amp_1, 'm--.', label = 'SPARC w/ outer code: after 1st round of AMP')
    ax.plot(EbN0_dB, BER_ldpc, 'm--', label = 'SPARC w/ outer code: after 1st round of LDPC')
    ax.plot(EbN0_dB, BER_amp_2, 'k-', label = 'SPARC w/ outer code: after 2nd round of AMP')
    if init=='soft':
        ax.plot(EbN0_dB, BER_ldpc_2, 'k--', label = 'SPARC w/ outer code: after 2nd round of LDPC')
    ax.plot(EbN0_dB, BER_plain, 'b-', label = 'Plain SPARC')
    if bpsk==True:
        ax.plot(EbN0_dB, BER_bpsk, 'g-', label = 'LDPC with BPSK modulation')
    plt.axvline(x=EbN0c_dB, color='r', linestyle='-', label='Shannon limit')
    plt.xlabel('$E_b/N_0$ (dB)', fontsize=15) # at some point need to work out how to write this so it outputs properly
    plt.ylabel('BER', fontsize=15)
    plt.tight_layout()
    plt.legend(loc=1, prop={'size': 7})
    plt.savefig(png_filename)


def soft_hard_plot(soft: bool, hard: bool, sec: int, soft_iter: int, sparcparams: SPARCParams, ldpcparams: LDPCParams, csv_filename: str, png_filename: str, datapoints=10, MIN_ERRORS=100, MAX_BLOCKS=500):
    """
    Plots soft information exchange and original hard information exchange.
    soft: determines if soft information exchange is performed
    hard: determines if original hard information exchange is performed. 
    sec: number of LDPC sections. Must be divisible by 8 to ensure nl is divisible by 24 (the number of columns in the proto matrix)
    soft_iter: number of iterations of soft information exchange. 
    Note need to change the plotting depending on the number of rounds you want to plot. 
    """
    # Sparc parameters
    L = sparcparams.L
    M = sparcparams.M
    logm = np.log2(M)
    p=sparcparams.p
    r_sparc = sparcparams.r
    T = sparcparams.t
    a = sparcparams.a
    f = sparcparams.f
    C = sparcparams.C 

    # LDPC parameters
    standard = ldpcparams.standard
    r_ldpc = ldpcparams.r_ldpc
    # number of ldpc bits, must be divisible by 24
    nl = logm * sec
    z = int(nl/24)
    ldpcparams = LDPCParams(standard, r_ldpc, z)

    n = L*logm/r_sparc
    R = (L*logm-nl*(1-5/6))/n

    print('Overall rate is: ', R)

    # snr at capacity
    snrc = (2**(2*R) - 1)
    EbN0c = 1/(2*R) * snrc
    EbN0c_dB = 20*np.log10(EbN0c)

    SIGMA = np.linspace(0.8, 0.4, datapoints)
    BER_sparc = np.zeros(datapoints)
    if soft:
        # Array of BER after ldpc decoding
        BER_ldpc_soft = np.zeros((datapoints, soft_iter))
        # Array of BER after AMP decoding
        BER_amp_soft = np.zeros((datapoints, soft_iter+1))
        
    if hard:
        BER_amp_hard = np.zeros((datapoints, 2))
        BER_ldpc_hard = np.zeros(datapoints)

    i=0
    for sigma in SIGMA:
        sparcparams = SPARCParams(L, M, sigma, p, r_sparc, T, a, f, C)
        # sparcparams for running at same overall rate but with no outer code
        sparcparams1 = SPARCParams(L, M, sigma, p, R, T, a, f, C)

        ber_sparc = np.zeros(MIN_ERRORS)
        for j in range(MIN_ERRORS):
            (ber_sparc[j],_,_,_) = amp_ldpc_sim(sparcparams1)
        BER_sparc[i] = np.sum(ber_sparc)/MIN_ERRORS

        if soft:
            # cumulative BER for amp and ldpc
            ber_cum_amp = np.zeros((soft_iter+1))
            ber_cum_ldpc = np.zeros(soft_iter)
            nblockerrors_ldpc = 0
            nblocks = 0
            while nblockerrors_ldpc < MIN_ERRORS:
                (ber_thisblock_amp, ber_thisblock_ldpc, _) = soft_amp_ldpc_sim(sparcparams, ldpcparams, soft_iter)
                #print("BER amp:", ber_thisblock_amp)
                #print("BER ldpc: ", ber_thisblock_ldpc

                ber_cum_amp += ber_thisblock_amp
                ber_cum_ldpc += ber_thisblock_ldpc
                if ber_thisblock_ldpc[0]:
                    nblockerrors_ldpc += 1
                nblocks += 1

                if nblocks >= MAX_BLOCKS:
                    break

            BER_amp_soft[i,:] = ber_cum_amp/nblocks
            BER_ldpc_soft[i,:] = ber_cum_ldpc/nblocks
        if hard:
            # cumulative BER for amp and ldpc
            ber_cum_amp = np.zeros(2)
            ber_cum_ldpc = 0
            ber_thisblock_amp = np.zeros(2)
            ber_thisblock_ldpc = 0
            nblockerrors_ldpc = 0
            nblocks = 0
            while nblockerrors_ldpc < MIN_ERRORS:
                (ber_thisblock_amp[0], ber_thisblock_ldpc, ber_thisblock_amp[1], _) = amp_ldpc_sim(sparcparams, ldpcparams)

                ber_cum_amp += ber_thisblock_amp
                ber_cum_ldpc += ber_thisblock_ldpc
                if ber_thisblock_ldpc:
                    nblockerrors_ldpc += 1
                nblocks += 1

                if nblocks >= MAX_BLOCKS:
                    break
            BER_amp_hard[i,:] = ber_cum_amp/nblocks
            BER_ldpc_hard[i] = ber_cum_ldpc/nblocks

        i+=1

    #snrdB = 20*np.log10(P/sigma**2)
    EbN0 = 1/(2*R) * (p/SIGMA**2)
    #print(EbN0)
    EbN0_dB = 20*np.log10(EbN0)
    # open file you want to write CSV output to. 'a' means its in append mode. Switching this to 'w' will make it overwrite the file.
    myFile = open(csv_filename, 'a')
    with myFile:
        if soft: 
            myFields = ['EbN0_dB', 'BER_sparc', 'BER_ldpc_soft', 'BER_amp_soft']
            writer = csv.DictWriter(myFile, fieldnames=myFields)
            writer.writeheader()
            for k in range(datapoints):
                writer.writerow({'EbN0_dB' : EbN0_dB[k], 'BER_sparc' : BER_sparc[k],'BER_ldpc_soft' : BER_ldpc_soft[k,:], 'BER_amp_soft': BER_amp_soft[k,:]})
        if hard: 
            myFields = ['EbN0_dB', 'BER_sparc', 'BER_ldpc_hard', 'BER_amp_hard']
            writer = csv.DictWriter(myFile, fieldnames=myFields)
            writer.writeheader()
            for k in range(datapoints):
                writer.writerow({'EbN0_dB' : EbN0_dB[k], 'BER_sparc' : BER_sparc[k],'BER_ldpc_hard' : BER_ldpc_hard[k], 'BER_amp_hard': BER_amp_hard[k,:]})


    fig, ax = plt.subplots()
    ax.set_yscale('log', basey=10)
    ax.plot(EbN0_dB, BER_sparc, 'g:', label = 'SPARC w/ same overall rate')
    if soft:
        # comment out and add lines below as appropriate
        ax.plot(EbN0_dB, BER_amp_soft[:,0], 'b:', label = 'SPARC w/ outer code: after AMP')
        ax.plot(EbN0_dB, BER_ldpc_soft[:,0], 'k:', label = 'SPARC w/ outer code: after LDPC (soft)')
        ax.plot(EbN0_dB, BER_amp_soft[:,1], 'b--', label = 'SPARC w/ outer code: after 2nd round of AMP (soft)')
        ax.plot(EbN0_dB, BER_ldpc_soft[:,1], 'k--', label = 'SPARC w/ outer code: after 2nd round of LDPC (soft)')
        ax.plot(EbN0_dB, BER_amp_soft[:,2], 'b-', label = 'SPARC w/ outer code: after 3rd round of AMP (soft)')
    if hard:
        #ax.plot(EbN0_dB, BER_amp_hard[:,0], 'c:', label = 'SPARC w/ outer code: after AMP (hard)')
        ax.plot(EbN0_dB, BER_ldpc_hard, 'm:', label = 'SPARC w/ outer code: after LDPC (hard)')
        ax.plot(EbN0_dB, BER_amp_hard[:,1], 'c--', label = 'SPARC w/ outer code: after 2nd round of AMP (hard)')
    plt.axvline(x=EbN0c_dB, color='r', linestyle='-', label='Shannon limit')
    plt.xlabel('$E_b/N_0$ (dB)') # at some point need to work out how to write this so it outputs properly
    plt.ylabel('BER')
    plt.tight_layout()
    plt.legend(loc=1, prop={'size': 7})
    plt.savefig(png_filename)


def soft_hardinit_plot(sparcparams: SPARCParams, ldpcparams: LDPCParams, csv_filename: str, png_filename: str, sections, datapoints=10, MIN_ERRORS=100, MAX_BLOCKS=500, soft_iter=3, threshold=0.6):
    """
    Plots threshold initialised information exchange. 
    sections: determines how many sections are covered by the LDPC code. Ensure it's compatible with z. 
    for standard 802.16, sections must be divisible by 8. This will be different for different ldpc codes.
    threshold: determines the threshold used for hard decoding. 
    soft_iter: the number of iterations of information exchange. 
    Produces one plot with results after the first and second iteration of info exchange and another plot with
    the results after the first and final iteration of info exchange. 
    """
    # Sparc parameters
    L = sparcparams.L
    M = sparcparams.M
    logm = np.log2(M)
    p=sparcparams.p
    r_sparc = sparcparams.r
    T = sparcparams.t

    # LDPC parameters
    standard = ldpcparams.standard
    r_ldpc = ldpcparams.r_ldpc
    z = ldpcparams.z
    if r_ldpc=='5/6':
        Rldpc=5/6
    elif r_ldpc=='1/2':
        Rldpc=1/2
    elif r_ldpc=='0.45':
        Rldpc=0.45
    elif r_ldpc=='3/8':
        Rldpc=3/8   
    else:
        print("Invalid choice of ldpc rate, please choose a different one.")

    nl = logm * sections
    # if I've set z externally, don't bother with this
    if z==None:
        if standard=="802.11n" or standard=="802.16":
            # no. ldpc bits must be divisible by 24 when using one of Jossy's codes
            z = int(nl/24)
        else:
            # number of ldpc bits, must be divisible by 40 when using my own ldpc
            z = int(nl/40)
    ldpcparams = LDPCParams(standard, r_ldpc, z)    

    n = L*logm/r_sparc
    R = (L*logm-nl*(1-Rldpc))/n
    # note that R should be r_ldpc
    print('Overall rate is: ', R)

    # snr at capacity -> found by inverting the equation C = 1/2*log2(1+snr) where C is the rate here
    snrc = (2**(2*R) - 1)
    EbN0c = 1/(2*R) * snrc
    EbN0c_dB = 20*np.log10(EbN0c)

    SIGMA = np.linspace(0.9, 1.4, datapoints)
    # BER after each round of AMP decoding
    BER_amp = np.zeros((datapoints, soft_iter))
    # BER after each round of ldpc
    BER_ldpc = np.zeros((datapoints, soft_iter))
    # BER of plain SPARC
    BER_plain = np.zeros(datapoints)

    i=0
    for sigma in SIGMA:
        sparcparams = SPARCParams(L, M, sigma, p, r_sparc, T)
        # plain sparc with same overall rate for comparison
        sparcparams_plain = SPARCParams(L, M, sigma, p, R, T)
        # cumulative BER for amp and ldpc
        ber_cum_amp = np.zeros(soft_iter)
        ber_cum_ldpc = np.zeros(soft_iter)
        ber_cum_plain = 0
        nblockerrors_plain = 0
        nblocks = 0
        while nblockerrors_plain < MIN_ERRORS:
            (ber_thisblock_amp, ber_thisblock_ldpc, _) = soft_amp_ldpc_hardinit(sparcparams, ldpcparams, soft_iter, threshold)
            #print("ber_thisblock_amp ", ber_thisblock_amp)
            #print("ber_thisblock_ldpc", ber_thisblock_ldpc)
            (ber_thisblock_plain, _, _, _) = amp_ldpc_sim(sparcparams_plain)
            for j in range(soft_iter):
                ber_cum_amp[j] += ber_thisblock_amp[j]
                ber_cum_ldpc[j] += ber_thisblock_ldpc[j]
            ber_cum_plain += ber_thisblock_plain
            if ber_thisblock_plain:
                nblockerrors_plain += 1
            nblocks += 1

            if nblocks >= MAX_BLOCKS:
                break

        BER_amp[i,:] = ber_cum_amp/nblocks
        BER_ldpc[i,:] = ber_cum_ldpc/nblocks
        #print("BER_ldpc ", BER_ldpc)
        BER_plain[i] = ber_cum_plain/nblocks
        i+=1

    #snrdB = 20*np.log10(P/sigma**2)
    EbN0 = 1/(2*R) * (p/SIGMA**2)
    #print(EbN0)
    EbN0_dB = 20*np.log10(EbN0)
    # open file you want to write CSV output to. 'a' means its in append mode. Switching this to 'w' will make it overwrite the file.
    myFile = open(csv_filename, 'a')
    with myFile:
        myFields = ['EbN0_dB', 'BER_amp', 'BER_ldpc', 'BER_plain']
        writer = csv.DictWriter(myFile, fieldnames=myFields)
        writer.writeheader()
        for k in range(datapoints):
            writer.writerow({'EbN0_dB' : EbN0_dB[k], 'BER_amp': BER_amp[k,:],'BER_ldpc' : BER_ldpc[k,:], 'BER_plain': BER_plain[k]})

    #print("BER_ldpc is: ", BER_ldpc)
    plt.figure(1)
    fig, ax = plt.subplots()
    ax.set_yscale('log', basey=10)
    # plot the results after the first and last soft iteration
    ax.plot(EbN0_dB, BER_amp[:,0], 'm--.', label = 'SPARC w/ outer code: after 1 round of AMP')
    ax.plot(EbN0_dB, BER_ldpc[:,0], 'm--', label = 'SPARC w/ outer code: after 1 round of LDPC')
    ax.plot(EbN0_dB, BER_amp[:,1], 'k-', label = 'SPARC w/ outer code: after 2nd round of AMP')
    ax.plot(EbN0_dB, BER_ldpc[:,1], 'k--.', label = 'SPARC w/ outer code: after 2nd round of LDPC')
    #ax.plot(EbN0_dB, BER_amp[:,soft_iter-1], 'k-', label = 'SPARC w/ outer code: after '+str(soft_iter)+' rounds of AMP')
    #ax.plot(EbN0_dB, BER_ldpc[:,soft_iter-1], 'k--.', label = 'SPARC w/ outer code: after '+str(soft_iter)+' rounds of LDPC')
    ax.plot(EbN0_dB, BER_plain, 'b-', label = 'Plain SPARC')

    plt.axvline(x=EbN0c_dB, color='r', linestyle='-', label='Shannon limit')
    plt.xlabel('$E_b/N_0$ (dB)', fontsize=15) # at some point need to work out how to write this so it outputs properly
    plt.ylabel('BER', fontsize=15)
    plt.tight_layout()
    plt.legend(loc=3, prop={'size': 7})
    plt.savefig(str('it2'+png_filename))


    plt.figure(2)
    fig, ax = plt.subplots()
    ax.set_yscale('log', basey=10)
    # plot the results after the first and last soft iteration
    ax.plot(EbN0_dB, BER_amp[:,0], 'm--.', label = 'SPARC w/ outer code: after 1 round of AMP')
    ax.plot(EbN0_dB, BER_ldpc[:,0], 'm--', label = 'SPARC w/ outer code: after 1 round of LDPC')
    #ax.plot(EbN0_dB, BER_amp[:,1], 'k--', label = 'SPARC w/ outer code: after 2nd round of AMP')
    #ax.plot(EbN0_dB, BER_ldpc[:,1], 'k-', label = 'SPARC w/ outer code: after 2nd round of LDPC')
    ax.plot(EbN0_dB, BER_amp[:,soft_iter-1], 'k-', label = 'SPARC w/ outer code: after '+str(soft_iter)+' rounds of AMP')
    ax.plot(EbN0_dB, BER_ldpc[:,soft_iter-1], 'k--.', label = 'SPARC w/ outer code: after '+str(soft_iter)+' rounds of LDPC')
    ax.plot(EbN0_dB, BER_plain, 'b-', label = 'Plain SPARC')

    plt.axvline(x=EbN0c_dB, color='r', linestyle='-', label='Shannon limit')
    plt.xlabel('$E_b/N_0$ (dB)', fontsize=15) # at some point need to work out how to write this so it outputs properly
    plt.ylabel('BER', fontsize=15)
    plt.tight_layout()
    plt.legend(loc=3, prop={'size': 7})
    plt.savefig(str('it'+ str(soft_iter)+ '_' +png_filename))


if __name__ == "__main__":
    # get the time so you can calculate the wall clock time of the process
    t0 = time.time()
    '''
    #####################################
    # testing out the soft information exchange with hard initialisation
    standard = '2_7_12_good'
    r_ldpc='1/2'
    z=32
    ldpcparams = LDPCParams(standard, r_ldpc, z)
    sparcparams = SPARCParams(L=256, M=32, sigma=None, p=4, r=1, t=64)

    soft_hardinit_plot(sparcparams, ldpcparams, csv_filename="testing.csv", png_filename="testing.png", datapoints=10, MIN_ERRORS=1, MAX_BLOCKS=1, soft_iter=2, threshold=0.7)
    print("Wall clock time elapsed: ", time.time()-t0)
    #csv_filename="shinit_M32L256Rsparc1P4_standardGood_dc6_Rldpc0_45z_32_it5_rep10.csv", png_filename="softhardinit_M32L256Ramp1P4_standardGood_dc6_Rldpc0_45z_32_it5_rep10.png",
    '''
    
    
    #####################################
    # Running new soft info exchange on original L=M=512 sparc with Jossy ldpc of rate 5/6
    # threshold initialisations info exchange
    # logm*sections needs to be divisible by 40. 
    #sections = 128 # gives overall rate 0.75
    #sections = 192 # gives overall rate 0.625
    #sections = 224 # gives overall rate 0.5625
    sections =256
    ldpcparams = LDPCParams('2_7_12_good', '1/2', z=None)
    sparcparams = SPARCParams(L=256, M=32, sigma=None, p=4, r=1, t=64)
    soft_hardinit_plot(sparcparams, ldpcparams, csv_filename="thresholdinit_M32L256Ramp1P4_std2_7_12_good_Rldpc1_2thres07_200reps_testPl.csv", png_filename="thresholdinit_M32L256Ramp1P4_std2_7_12_good_Rldpc1_2thres07_200reps_testPl.pdf", datapoints=10, MIN_ERRORS=200, MAX_BLOCKS=250, soft_iter=3, threshold=0.7, sections=sections)
    
    #ldpcparams2 = LDPCParams('2_7_12_bad', '1/2', z=None)
    #soft_hardinit_plot(sparcparams, ldpcparams2, csv_filename="thresholdinit_M32L256Ramp1P4_std2_7_12_bad_Rldpc1_2z32thres07_300reps.csv", png_filename="thresholdinit_M32L256Ramp1P4_std2_7_12_bad_Rldpc1_2z32thres07_300reps.pdf", datapoints=10, MIN_ERRORS=300, MAX_BLOCKS=350, soft_iter=4, threshold=0.7, sections=sections)
    
    print("Wall clock time elapsed: ", time.time()-t0)    

    '''
    ######################################
    # Plot the parameterised power allocation
    datapoints = 5
    P=2.5
    L=512
    R=1
    SNR_dB = np.linspace(10, 14, datapoints)
    SNR = 10**(SNR_dB/20)
    pa = np.zeros((datapoints,L))
    i=0
    for snr in SNR:
        # The channel capacity of the AWGN channel 
        #snr = P / sigma**2
        C = 0.5 * np.log2(1 + snr)
        # A good starting point for the parameters is:
        a=R/C
        f=R/C
        pa[i,:] = pa_parameterised(L, C, P, a, f)
        i=i+1
    # checking that the total power adds up to P
    print('All the following numbers should be '+ str(P) + ': ', np.sum(pa,1))
    fig, ax = plt.subplots()
    ax.plot(pa[0,:], 'b--', label='$SNR_{dB}$='+str(SNR_dB[0])) 
    ax.plot(pa[1,:], 'k--', label='$SNR_{dB}$='+str(SNR_dB[1]))
    ax.plot(pa[2,:], 'm--', label='$SNR_{dB}$='+str(SNR_dB[2]))
    ax.plot(pa[3,:], 'c--', label='$SNR_{dB}$='+str(SNR_dB[3]))
    ax.plot(pa[4,:], 'r--', label='$SNR_{dB}$='+str(SNR_dB[4]))

    plt.xlabel('L')
    plt.ylabel('Power Allocation')
    plt.legend(loc=6, prop={'size': 7})
    plt.title("The power allocation for a rate 1 SPARC with different SNRs")
    plt.show()
    '''

    
    '''
    ##########################################
    # Plot waterfall curves
    # Compare ldpc with bpsk of rate 5/6 to a sparc with sparc rate 1 and ldpc rate 5/6 with all sections covered
    # to give a overall rate of 5/6
    # And SPARC with no outer code and overall rate 5/6
    # Note that z is set within the waterfall function so just set as None here
    # sections for use with originalHard.
    sections = 384 
    ldpcparams = LDPCParams('802.16', '5/6', None)
    sparcparams = SPARCParams(L=512, M=512, sigma=None, p=4, r=1, t=64)
    waterfall(sparcparams, ldpcparams, init='originalHard', pa_param=False, datapoints=10, MIN_ERRORS=200, MAX_BLOCKS=250, csv_filename='test.csv', png_filename='test.pdf', sections=sections)

    print("Wall clock time elapsed: ", time.time()-t0)
    '''
    '''
    #########################################
    # Plot hard and soft loops
    ldpcparams = LDPCParams('802.16', '5/6', None)
    sparcparams = SPARCParams(L=768, M=512, sigma=None, p=1.8, r=1, t=64)
    soft_hard_plot(soft=True, hard=True, sec=569, soft_iter=2, sparcparams= sparcparams, ldpcparams=ldpcparams, csv_filename='EbN0VsBER_soft_hard_100_4.csv', png_filename='EbN0VsBER_soft_hard_100_4.png', datapoints=10, MIN_ERRORS=100, MAX_BLOCKS=100)
    
    print("Wall clock time elapsed: ", time.time()-t0)
    '''
    