#!/usr/bin/env python3
from perseus.client import ModuleClient
import matplotlib.pyplot as pl
import matplotlib as mpl
import sys, time, scipy.signal
import os
import sys
import numpy as np

# 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.

def main(args):
    m = ModuleClient(args.dk_name)
    m.tuber_resolve()

    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): ', 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): ', t.trigger_rates(3))

    wf = [np.asarray(m.trigger.get_waveforms(100000)) for i in range(args.n_samples)]
    # wf = m.trigger.get_waveforms(100000)
    wf = np.hstack(wf)

    fig, p = pl.subplots(2, 2, figsize=(18, 12))
    cm = pl.colormaps['cividis']
    cm2 = pl.colormaps['vanimo']

    # Spectral Density plots always in NE plot
    for i in range(16):
        f, pspec = scipy.signal.periodogram(wf[i] - np.mean(wf[i]), fs=1./4.8e-9, return_onesided=True, scaling='spectrum', window='blackman')
        pspec /= 2*2048**2
        fbins = np.logspace(4.5,8.5, 80)
        rebinned = np.histogram(f, weights=pspec, bins=fbins)
        pspec_dbfs = 10*np.log10(rebinned[0]/np.diff(rebinned[1]))

        if i < 8 :
            p[0,1].plot((rebinned[1][1:] + rebinned[1][:-1])/2, pspec_dbfs, '-.',  linewidth=1, alpha=0.9, label='Channel %d' % i, color=cm(i/7.))
        else:
            p[0,1].plot((rebinned[1][1:] + rebinned[1][:-1])/2, pspec_dbfs, '--',  linewidth=1, alpha=0.9, label='Channel %d' % i, color=cm2((i-7)/7.))

    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)

    # Channel Correlation plots
    if args.corr == False:

        # Show the channel waveforms in NW plot
        for i in range(16):

            if i < 8 :
                p[0,0].plot(wf[i], '-.', label='Channel %d' % i, linewidth=0.5, alpha=0.9, color=cm(i/7.))
            else:
                p[0,0].plot(wf[i], '--', label='Channel %d' % i, linewidth=0.5, alpha=0.9, color=cm((i-7)/7.))

        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)

        # Show the channel correlation moments (coefficients) in SW plot
        mat = p[1,0].imshow(np.corrcoef(wf))
        p[1,0].set_title('Channel-to-channel Correlation Coefficients')
        fig.colorbar(mat)

        # Show the total noise per channel in SE plot
        p[1,1].plot(np.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)')
    else:

        # Show the channel correlation moments (coefficients) in NW plot
        mat = p[0,0].imshow(np.corrcoef(wf))
        p[0,0].set_title('Channel-to-channel Correlation Coefficients')
        fig.colorbar(mat)

        # Calculate a statistical measure of signal injection into other channels
        tens = np.zeros((16, 16))
        for i in range(len(tens)):
            row_mean = np.mean(wf[i])
            dot_prod = np.dot(wf[i]-row_mean, wf[i]-row_mean)
            tens[i,:] = np.dot(wf[:] - np.mean(wf[:]), wf[i]-row_mean)/dot_prod

        print(tens)
        int_noise = np.std(wf, axis=1)
        # Discern numerically which channels are getting an analog signal.
        # SW plots are filled with signal leakage
        # SE plots are filled with log_10(abs(signal)) leakage
        if np.mean(int_noise[0:8]) >= np.mean(int_noise[8:]):
            mat2 = p[1,0].imshow(tens[0:8])
            p[1,0].set_ylabel("Agressor Channels: CH0 - CH7")
            fig.suptitle("Analog test with CH 0-7 analog input", fontsize=20)

            mat1 = p[1,1].imshow(np.log10(np.abs(tens[0:8])))
            p[1,1].set_ylabel("Agressor Channels: CH0 - CH7")

        else:
            mat2 = p[1,0].imshow(tens[8:])
            p[1,0].set_ylabel("Agressor Channels: CH7 - CH15")
            fig.suptitle("Analog test with CH 8-15 analog input", fontsize=20)

            mat1 = p[1,1].imshow(np.log10(np.abs(tens[8:])))
            p[1,1].set_ylabel("Agressor Channels: CH7 - CH15")

            p[1,0].set_yticklabels([i for i in range(7,16,1)])
            p[1,1].set_yticklabels([i for i in range(7,16,1)])

        fig.colorbar(mat1)
        fig.colorbar(mat2)
        p[1,1].set_title(r'$Log_{10} (\langle A \cdot B\rangle / \langle A \cdot A\rangle)$')
        p[1,0].set_title(r'$\langle A \cdot B\rangle / \langle A \cdot A\rangle$')

    print(f"Saved file to\n\t-> {args.data_out}")
    fig.savefig(f"{args.data_out}")

if __name__ == "__main__":
    import argparse

    p = argparse.ArgumentParser()

    p.add_argument("--dk_name", "-name", type=str, required=True,
                   help="Database unique name for mainboard. Also referred to the IP name. You DO NOT need to include the local suffix")
    p.add_argument('--data_out', '-out', type=str, required=True,
                   help="System path to save the output figure from this test.")
    p.add_argument('--corr', '-cr', type=bool, required=False, default=False,
                   help="Change some of the studies for a correlation study between channels.")
    p.add_argument('--n_samples', '-N', type=int, required=False, default=20,
                   help='Number of 100k samples to take from ADC')

    p.set_defaults(func=main)

    args = p.parse_args()

    args.dk_name = f"{args.dk_name}.local"

    args.func(args)