#!/usr/bin/env python

# Quick and dirty script to plot some Biphase Mark Code explanation figures.
# 2011, M. Elizabeth Scott <beth@scanlime.org>

import numpy
from matplotlib import pyplot

def biphase_mark_encode(bits):
    out = []
    state = 0
    for bit in bits:
        state = state ^ 1
        out.append(state)
        state = state ^ bit
        out.append(state)
    return out

def plot_bitstream_time(axis, bits, numSamples=10):
    samples = numpy.repeat(bits, numSamples)
    axis.plot([i / float(numSamples) + 0.5 for i in range(len(samples))], samples)
    axis.set_xlim([0.5, len(bits)+0.5])
    axis.set_ylim([-0.5, 1.5])
    axis.get_xaxis().set_ticks(range(1, 1+len(bits)))
    axis.get_yaxis().set_visible(False)

def plot_bitstream_freq(axis, bits, bitRate, unit, numSamples=50):
    samples = numpy.repeat(bits, numSamples)
    sampleRate = bitRate * numSamples
    unit /= 1e6

    axis.plot(numpy.fft.fftfreq(len(samples), d=1.0/sampleRate) / 1e6,
              abs(numpy.fft.fft(samples)) / len(samples) / 0.015)
    axis.set_xlim([0, unit*3])
    axis.set_ylim([0, 1])
    axis.get_xaxis().set_ticks([0, unit, unit*2])
    axis.get_yaxis().set_visible(False)

# Bitstream sample

sample = (0,0,0,1,1,1,0,1,0,1)

fig = pyplot.figure(figsize=(10, 4))
ax = fig.add_subplot(211)
ax.set_title("Sample Bitstream, Before and After BMC")
plot_bitstream_time(ax, sample)
ax = fig.add_subplot(212)
plot_bitstream_time(ax, biphase_mark_encode(sample))

fig.savefig('bitstream-sample.png')

# Spectrum comparison

fig = pyplot.figure(figsize=(10, 6))
noise = numpy.random.random_integers(0, 1, 10000)
rate = 48000 * 64

ax = fig.add_subplot(211)
ax.set_title("Spectrum of 64x 48 KHz Random Bitstream, Before and After BMC")

plot_bitstream_freq(ax, noise, rate, rate)
ax = fig.add_subplot(212)
ax.set_xlabel("Frequency (MHz)")
plot_bitstream_freq(ax, biphase_mark_encode(noise), rate*2, rate)

fig.savefig('spectrum-comparison.png')