#!/usr/bin/env python3
from perseus.client import ModuleClient

import sys, time, numpy, scipy.signal
import os
import sys

# Run in a pylab environment. Does not save plots, but will show them.
# Run e.g. as ipython --pylab -i waveform-stats.py discokraken_XXXXXX.local
#
# Shows a snapshot of the waveforms, power spectra, cross-correlation, and
# noise statistics. The trigger rates shown probe non-Gaussianity beyond
# the quoted noise levels.
# 
# An optional set of channels to plot can follow the hostname, e.g.
# ipython --pylab -i waveform-stats.py discokraken_XXXXXX.local 0 3 6
# will show data for channels 0, 3, and 6 only

m = ModuleClient(sys.argv[1])
m.tuber_resolve()

channels = list(range(16))
if len(sys.argv) > 2:
    channels = [int(i) for i in sys.argv[2:]]

t = m.trigger
print('Running for %.2f seconds' % (t.current_sample()/208.3333333e6))
print('Mean: ', t.channel_means(100))
print('RMS: ', t.channel_rms(100))
print(t.timecode_sync_data())

t.set_thresholds([int(i + 3) for i in t.channel_means()])
t.enable_scalers()
print('Rates at +3 counts trigger (Hz): ', [('%.1f' % i) for i in t.trigger_rates(3)])

t.set_thresholds([int(i + 5) for i in t.channel_means()])
t.enable_scalers()
print('Rates at +5 counts trigger (Hz): ', [('%.1f' % i) for i in t.trigger_rates(3)])

wf = numpy.asarray(m.trigger.get_waveforms(100000))

fig, p = pylab.subplots(2, 2)

for i in channels:
    p[0,0].plot(wf[i], label='Channel %d' % i)
p[0,0].set_xlabel('Sample Index (4.8 ns)')
p[0,0].set_ylabel('ADC Counts')
p[0,0].set_title('Sample Waveforms')
p[0,0].legend(ncols=4)

for i in channels:
    f, pspec = scipy.signal.periodogram(wf[i] - numpy.mean(wf[i]), fs=1./4.8e-9, return_onesided=True, scaling='spectrum', window='blackman')
    pspec /= 2*2048**2
    fbins = numpy.logspace(4.5,8.5, 200)
    rebinned = numpy.histogram(f, weights=pspec, bins=fbins)
    pspec_dbfs = 10*numpy.log10(rebinned[0]/numpy.diff(rebinned[1]))
    p[0,1].plot((rebinned[1][1:] + rebinned[1][:-1])/2, pspec_dbfs, label='Channel %d' % i)
p[0,1].set_xlabel('Frequency (Hz)')
p[0,1].set_ylabel('Spectral Density (dBFS/Hz)')
p[0,1].semilogx()
p[0,1].grid()
p[0,1].set_ylim(-180, None)
p[0,1].legend(ncols=4)

mat = p[1,0].imshow(numpy.corrcoef(wf))
p[1,0].set_title('Channel-to-channel Correlation Coefficients')
fig.colorbar(mat)

p[1,1].plot(numpy.std(wf, axis=1))
p[1,1].set_title('Noise vs. Channel Number')
p[1,1].set_xlabel('Channel')
p[1,1].set_ylabel('Noise (counts)')

fig.show()