# This is the template file for Lab #4, Task #5
import numpy,random,math
import lab4
import matplotlib.pyplot as p
from lab4_2 import *
from lab4_3 import ViterbiDecoder
from lab4_4 import SoftViterbiDecoder

# This program takes a while to run (45 minutes on my Macbook air 1.86
# GHz dual core computer, but using only one core).  Please be patient
# -- it'll be worth your while!
if __name__=='__main__':
    # Various arrays structures used to produce plots at the end.
    nlist = []
    ber_orig = []
    ber_rect = []
    ber_hard = []
    ber_hard4 = []
    ber_soft = []
    ber_soft4 = []
    E_b = 0.25               # energy/bit w/ 0V for "0" and 1V for "1"

    nbits = 200000
    K3 = 3; glist3 = (7,6)
    K4 = 4; glist4 = (0xE,0xD)

    if nbits > 20000:
        print 'This program will take a while.  Don\'t hold your breath.  Go make yourself some coffee or dinner!'
        print

    # Make up a random message, run it through the two convolutional
    # encoders, Do the same for the stream block code with rectangular
    # parity.
    message = numpy.random.random_integers(0,1,nbits)
    encoded3 = lab4.convolutional_encoder(message, K3, glist3)
    encoded4 = lab4.convolutional_encoder(message, K4, glist4)
    cw = []
    for i in range(0,nbits,4): cw.extend(lab4.codeword(message[i:i+4], 2, 2))

    for d in range(7,-1,-1):
        db = 0.5*d              # decibel scale; E_b/N_0; N_0/2 = sigma^2
        sigma = math.sqrt(E_b/(2*math.pow(10, db/10.0)))
        print "nbits: %d; Eb/N0 %.1f; noise (2*variance): %7.6f" \
            % (nbits, db, 2*sigma*sigma)
        nlist.append(db)
        noisy = numpy.random.normal(0,sigma,2*len(message))
        received = encoded3 + noisy 
        received4 = encoded4 + noisy 
            
        noisy_bits = (received >= 0.5)
        BER_orig = lab4.ber(encoded3, noisy_bits)
        print "  BER of original reception =  %7.6f" % BER_orig
        ber_orig.append(BER_orig)

        # Try out the rectangular parity code with rate 1/2
        rcd =  (cw + noisy) >= 0.5
        rect_decoded = decode_stream(rcd, 2, 2)
        BER_rect = lab4.ber(message, rect_decoded)
        print "  BER using rect parity = %7.6f" % BER_rect
        ber_rect.append(BER_rect)

        # Try out hard decision Viterbi decoder of constraint length 3
        hard = ViterbiDecoder(K3, glist3)
        received_hard = hard.decode(received)
        BER_hard = lab4.ber(message,received_hard)
        print "  BER using hard-decision branch metric (K=3) = %7.6f" % BER_hard
        ber_hard.append(BER_hard)

        # Try out hard decision Viterbi decoder of constraint length 4
        hard4 = ViterbiDecoder(K4, glist4)
        received_hard = hard4.decode(received4)
        BER_hard4 = lab4.ber(message,received_hard)
        print "  BER using hard-decision branch metric (K=4) = %7.6f"%BER_hard4
        ber_hard4.append(BER_hard4)

        # Try out soft decision Viterbi decoder of constraint length 3
        soft = SoftViterbiDecoder(K3, glist3)
        received_soft = soft.decode(received)
        BER_soft = lab4.ber(message,received_soft)
        print "  BER using soft-decision branch metric (K=3) = %7.6f" % BER_soft
        ber_soft.append(BER_soft)

        # Try out soft decision Viterbi decoder of constraint length 4
        soft4 = SoftViterbiDecoder(K4, glist4)
        received_soft = soft4.decode(received4)
        BER_soft4 = lab4.ber(message,received_soft)
        print "  BER using soft-decision branch metric (K=4) = %7.6f"%BER_soft4
        ber_soft4.append(BER_soft4)

        nbits = max(20000,0.8*nbits)
        
    p.xlabel('Energy_per_bit/Noise')
    p.ylabel('Prob. of decoding error (BER)')
    p.semilogy(nlist, ber_orig, color='green', marker='o', label="uncoded", linewidth=2)
    p.semilogy(nlist, ber_rect, 'cyan', marker='o', label="rect parity", linewidth=2)
    p.semilogy(nlist, ber_hard, 'red', marker='o', label="K=3 hard Viterbi", linewidth=2)
    p.semilogy(nlist, ber_hard4, 'magenta', marker='o', label="K=4 hard Viterbi", linewidth=2)
    p.semilogy(nlist, ber_soft, 'bo-', marker='o', label="K=3 soft Viterbi", linewidth=2)
    p.semilogy(nlist, ber_soft4, 'black', marker='o', label="K=4 soft Viterbi", linewidth=2)
    p.legend(numpoints=1, loc='lower left')
    p.savefig("lab4data.ps", format='ps')
    p.show()