import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

x = np.arange(500)
y = 0.3 + np.sin(x/20.)/5
noisy = y + np.random.rand(500)/3 - (3./20)

plt.plot(x, y)
plt.plot(x, noisy, alpha=0.5)
plt.show()

plt.figure(figsize=(12,4))

plt.subplot(121)
plt.plot(x, y, lw=2, color='k')
plt.text(10, 0.05, "signal we'd like to detect", color='k')
plt.hlines(0.6, 0, 500, color='b', linestyles='--')
plt.text(10, 0.62, 'detection threshold', color='b')
plt.ylabel('amplitude')
plt.xlabel('time (ms)')
plt.ylim(0, 0.8)
plt.xlim(0, 500)

plt.subplot(122)
plt.plot(x, noisy, color='g', alpha=0.8)
plt.plot(x, y, color='k', lw=2)
plt.hlines(0.6, 0, 500, color='b', linestyles='--')
plt.text(10, 0.69, "additive noise", color='g')
plt.xlabel('time (ms)')
plt.ylim(0, 0.8)
plt.xlim(0, 500)

plt.show()

signal_level = 0.7
noise_level = 1.0
threshold_level = 1.1

im = np.zeros((64,64))
im[16:48,16:48] = signal_level

plt.imshow(im, interpolation="None")
plt.set_cmap('Greys')
plt.colorbar()
plt.show()

def threshold(im, value=0.5, mode='same'):
    t = np.copy(im)
    t[t<value] = 0
    if mode == 'binary':
        t[t>=value] = 1
    return t

plt.imshow(threshold(im, value=threshold_level), interpolation="None")
plt.show()

noise = np.random.rand(64,64) * noise_level - noise_level/2
noised = im + noise
plt.imshow(noised, interpolation="None")
plt.colorbar()
plt.show()

noise = np.random.rand(64,64) * noise_level - noise_level/2
noised = im + noise
thresholded = threshold(noised, threshold_level)
plt.imshow(thresholded, interpolation="None", vmin=1)
plt.colorbar()
plt.show()

plt.figure(figsize=(18,4))

plt.subplot(141)
plt.title('Signal at {}'.format(signal_level))
plt.imshow(im, interpolation="None", cmap="Greys")
plt.ylim(0, 64)
plt.xlim(0, 64)
plt.hlines(32, 0, 64, color='r')

plt.subplot(142)
plt.title('Noise at {0}'.format(noise_level))
plt.imshow(noise, interpolation="None", cmap="Greys")
plt.ylim(0, 64)
plt.xlim(0, 64)
plt.hlines(32, 0, 64, color='r')

plt.subplot(143)
plt.title('Section across signal and noise', color='r')
plt.plot(im[32,:], lw=2, color='k')
plt.plot(noised[32,:], color='g')
plt.ylim(-0.1, 1.3)
plt.xlim(0, 64)
plt.hlines(1.1, 0, 64, color='b')
plt.text(3, 1.12, 'threshold', color='b')
plt.text(6, 0.63, 'signal', color='k')
plt.text(15, 0.90, 'noisy\nsignal', color='g', horizontalalignment='right')

plt.subplot(144)
plt.title('Detected signal')
plt.imshow(thresholded, interpolation="None", cmap="Greys", vmin=1)
plt.hlines(32, 0, 64, color='r')
plt.show()

plt.show()

rec_noise = np.random.rand(64,64) + 0.5
recorded = thresholded + rec_noise

plt.imshow(recorded, interpolation="None", cmap="Greys", vmin=1)
plt.show()

import scipy.stats

signal_level = 0.3
noise_level = 0.4
threshold_level = 1

im = np.zeros((64,64))
im[16:48,16:48] = signal_level

# Use a different distribution, with a nasty hack for the standard deviation
noise = np.random.lognormal(noise_level, (noise_level+1)/6, 64*64).reshape((64,64))
#noise = np.random.rand(64,64) * noise_level - noise_level/2
noised = im + noise

thresholded = threshold(noised, threshold_level)

plt.figure(figsize=(18,5))

plt.subplot(131)
plt.imshow(threshold(im, threshold_level), interpolation="None")
plt.title('S:N {0}'.format(scipy.stats.signaltonoise(threshold(im, threshold_level), axis=None)))
plt.colorbar(shrink=0.8)

plt.subplot(132)
plt.imshow(noised, interpolation="None")
plt.title('S:N {0}'.format(scipy.stats.signaltonoise(noised, axis=None)))
plt.colorbar(shrink=0.8)

plt.subplot(133)
plt.imshow(thresholded, interpolation="None")
plt.title('S:N {0}'.format(scipy.stats.signaltonoise(thresholded, axis=None)))
plt.colorbar(shrink=0.8)
plt.show()

signal_level = 0.3
threshold_level = 1

im = np.zeros((64,64))
im[16:48,16:48] = signal_level

ston = []
r = np.arange(0,3.,0.05)

for n in r:
    noise = np.random.lognormal(n, (n+1)/6, 64*64).reshape((64,64))
    #noise = np.random.rand(64,64) * n - n/2
    noised = im + noise
    thresholded = threshold(noised, threshold_level)
    ston.append(scipy.stats.signaltonoise(thresholded, axis=None))
    
plt.plot(r, ston)
plt.title('Max at {0}'.format(r[max(enumerate(ston),key=lambda x: x[1])[0]]))
plt.ylabel('signal:noise')
plt.xlabel('mean of noise distribution')
plt.show()