%pylab inline
import seaborn
style.use('seaborn-poster')
rcParams['lines.linewidth'] = .5

#bits = where(random.random(32) > .5, 0, 1)
bits = array([0, 1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 0, 0, 1, 0, 0,
              0, 0, 0, 1, 1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0])
bits

fs = int(1e6)  # sampling frequency
baud = 9600
bit_time = fs / baud   # samples per bit
jitter = -31           # first bit starts 31 samples before we start sampling
measurement_time = int(ceil(bit_time * (len(bits) - 1)))
t = arange(measurement_time) / float(fs)
bit_index = ((t - jitter / fs) * baud).round().astype(int)
a = where(bits[bit_index], 5, 0)
subplot(411); ylim(-1, 6); plot(t, a); ylabel('V'); xlabel('s'); title('a: original pristine NRZ binary signal')
bit_time

b = a + random.normal(0, 10, size=a.shape)
subplot(411); plot(t, b); ylabel('V'); xlabel('s'); title('b: sadly corrupted with noise :\'(')
(a**2).mean()**0.5, ((b-a)**2).mean()**0.5

pulse = ones(int(floor(bit_time)))                 # the pulse shape of a single bit
c = convolve(b, pulse, mode='full') / sum(pulse)   # output is N+M-1 samples. this is not the fastest way to box filter
ct = (arange(len(c)) - len(pulse)//2) / float(fs)  # c time
subplot(411); plot(ct, c); xlabel('s'); title('c: box-filtered to remove noise while minimally slowing transitions')
len(c)

d = abs(c - c.mean())
subplot(411); plot(ct, d); xlabel('s'); title('d: distance of filtered signal from its mean')
d

n_impulses = 16
impulses = zeros(int(ceil(n_impulses * bit_time)))
impulses[(arange(n_impulses) * bit_time).round().astype(int)] = 1
it = arange(len(impulses)) / float(fs)
subplot(211); plot(it, impulses, label="impulse train at bit transition frequency"); legend(); ylim(-1, 2); tick_params(labelbottom=0)
e = convolve(d, impulses, mode='full') / impulses.sum()
et = (arange(len(e)) - len(pulse)//2 - len(impulses)//2) / float(fs)
ewr = int(round(bit_time / 3))   # e window radius
e_is_max = array([(e[i] >= e[i-ewr:i+ewr]).all() for i in range(len(e))])

subplot(212, sharex=gca()); plot(et, e*5 + 2, label="e: d convolved with impulse train (scaled and filtered)")
axhline(c.mean(), color="#ff6699"); xlim(t.min(), t.max()); plot(ct, c, label="c: filtered noisy signal")

for ti in et[e_is_max]:
    if t.min() <= ti <= t.max():
        axvline(ti, color="#ff99cc")
        
legend(loc=0); xlabel('s')
e_is_max.sum()

subplot(411); plot(et, e_is_max); plot(et, e)

f = []
ft = []
sample_times = set(ti for ti in et[e_is_max] if t.min() <= ti <= t.max())
for ci, cti in zip(c, ct):
    if cti in sample_times:
        f.append(ci)
        ft.append(cti)
        
f = array(f)
        
subplot(411); stem(ft, f); xlim(t.min(), t.max()); plot(ct, c); axhline(c.mean(), color="#ff6699"); ylabel('V')
subplot(412, sharex=gca()); plot(t, a); ylim(-1, 6)
f.round(2)

subplot(411); stem(ft, f > c.mean()); ylim(-1, 2)
subplot(412, sharex=gca()); plot(t, a); ylim(-1, 6)

where(f > c.mean(), 1, 0), bits