import numpy
import channel
reload(channel)
import matplotlib.pyplot as p
import lab1
reload(lab1)
import lab2
reload(lab2)

p.ion()  

def all_bit_patterns(n):
    # generate all possible n-bit bit sequences and merge into a single array
    message = [0] * (n * 2**n)
    for i in xrange(2**n):
        for j in xrange(n):
            message[i*n + j] = 1 if (i & (2**j)) else 0
    return numpy.array(message,dtype=numpy.int)


def plot_eye_diagram_fast(channel,plot_label, samples_per_bit):
    """
    Plot eye diagram for given channel using all possible 6-bit patterns
    merged together as message. plot_label is the label used in the eye
    diagram, Samples_per_bit determines how many samples
    are sent for each message bit.
    """

    # build message
    message = all_bit_patterns(6)
    
    # send it through the channel
    result = channel(lab1.bits_to_samples(message,samples_per_bit))

    # Truncate the result so length is divisible by 3*samples_per_bit
    result = result[0:len(result)-(len(result) % (3*samples_per_bit))]

    # Turn the result in to a n by 3*samples_per_bit matrix, and plot each row
    mat_samples = numpy.reshape(result,(-1,3*samples_per_bit))
    p.figure()
    p.plot(mat_samples.T)
    p.title('Eye diagram for channel %s' % plot_label)


def plot_eye_diagram(channel,plot_label, samples_per_bit):
    """
    Your, more educational, version of plot_eye_diagram.
    For given channel, you should generate a bit sequence made of
    all possible 3 bit messages (you can use the all_bit_patterns function
    above), and then plot the received samples by overlaying sets of
    3*samples_per_bit.  So that you can watch the eye diagram form,
    have the program stop for a press of enter after every 3*samples_per_bit
    plot The python function raw_input('Press enter to continue') will help.
    """

    pass # Your code here.


if __name__ == '__main__':

    # Create the channels (noise free with no random delays or padding)
    channel0 = channel.channel(channelid='0')
    channel1 = channel.channel(channelid='1')
    channel2 = channel.channel(channelid='2')

    # plot the eye diagram for the three virtual channels
    plot_eye_diagram(channel0,'0',samples_per_bit=50)
    plot_eye_diagram(channel1,'1',samples_per_bit=50)
    plot_eye_diagram(channel2,'2',samples_per_bit=50)

    # Create the channels with noise
    channel0 = channel.channel(channelid='0', noise=0.05)
    channel1 = channel.channel(channelid='1', noise=0.05)
    channel2 = channel.channel(channelid='2', noise=0.05)

    # plot the eye diagram for the three virtual channels
    plot_eye_diagram(channel0,'0n',samples_per_bit=50)
    plot_eye_diagram(channel1,'1n',samples_per_bit=50)
    plot_eye_diagram(channel2,'2n',samples_per_bit=50)