# template for Lab #3, Task #6
import numpy,random
import matplotlib.pyplot as p
import channel
reload(channel)
import lab1
reload(lab1)
import lab2
reload(lab2)
import lab2_1
reload(lab2_1)
import lab2_2
reload(lab2_2)
import lab2_4
reload(lab2_4)
import lab2_5
reload(lab2_5)
import lab2_6
reload(lab2_6)
import lab3_4
reload(lab3_4)

# turn on interactive mode, useful if we're using ipython
p.ion()  


# THIS IS THE FUNCTION YOU SHOULD MODIFY!
# This function should take a numpy array unit sample response hin,
# modify it so that deconvolution is less sensitive to noise, and then
# return the resulting unit sample response as a numpy array
def modify_h(hin):
    # Make a copy of hin and then modify hout
    hout = numpy.copy(hin)
    # Your program here, modify hout********************************

    # Helpful debugging plots
    p.figure()
    p.stem(range(len(hin)), hin)
    p.stem(range(len(hout)), hout)
    p.title('Unmodified and modified unit sample responses')
    
    return hout

# Modified deconvolved class from Lab2 Task 6.
class deconvolved_channel:
    def __init__(self, orig_channel, avgs=10, modifyit=True):
        self.orig_channel = orig_channel

        usr = lab2_2.unit_sample_response(orig_channel,max_length=400,tol=0.0)
        for i in range(avgs-1):
            usr += lab2_2.unit_sample_response(orig_channel,max_length=400,tol=0.0)
        usr /= avgs
        p.figure()
        p.stem(range(len(usr)), usr)
        if modifyit: 
            self.h = modify_h(usr) # Insures usr not overwritten

        p.figure()
        p.stem(range(len(self.h)), self.h)

    def __call__(self, input):
        h = self.h
        rcvd = self.orig_channel(numpy.append(input,numpy.zeros(len(h))))
        deconv = lab2_5.deconvolver(rcvd,h)
        return deconv[0:len(input)]


if __name__ == '__main__':

    # Create scaled_ir channel
    myir = lab2_4.scale_ir()

    # Deconvolved channel.
    dchannelir = deconvolved_channel(myir,modifyit=True)

    # Okay, let er rip.
    samples_per_bit = 30

    lab2_1.plot_eye_diagram_fast(myir,'ir', samples_per_bit=samples_per_bit)
    lab2_1.plot_eye_diagram_fast(dchannelir,'post-processed ir', samples_per_bit=samples_per_bit)

    pad_bits = 10
    sync_seq = [1,1,0,0,0,0,0,1,0,1]

    # Generate 500 random bits
    bit_seq = [random.randint(0,1) for i in xrange(500)]

    # Generate bit sequence with a start sync, plus lead and trail padding.
    # Turn sequences into numpy arrays for ease of processing later.
    sbits = numpy.array([0]*pad_bits+[1]*pad_bits+sync_seq+bit_seq+[0]*pad_bits)
    sync = numpy.array(sync_seq)
    bits = numpy.array(bit_seq)

    # Send bits through the channel, receive, remove sync and return bits.
    rcvd = lab3_4.xmit_and_rcv(sbits, dchannelir, samples_per_bit, sync)

    # Eliminate trailing bits
    rcvd = rcvd[0:len(bits)]

    # Check that enough bits were received
    lr = len(rcvd)
    ls = len(bits)
    assert lr==ls,"oops, rcvd bits=%d, sent bits=%d" % (lr,ls)


    # Finally, compute the bit error rate with trailing rcvd bits removed.
    ber = lab3_4.bit_error_rate(bits, rcvd)

    print "bit error rate = %g" % ber