%pylab inline
import calin.ix.scripts.compute_diagnostics
import calin.io.sql_transceiver
import calin.diagnostics.waveform
import calin.diagnostics.functional
from calin.plotting import plot_histogram
from matplotlib.backends.backend_pdf import PdfPages
from cycler import cycler

"""
Explanation: Plot the channel diagnostics
calin/examples/iact_data/plot channel diagnostics.ipynb - Stephen Fegan - 2017-03-27
Copyright 2017, Stephen Fegan &#115;&#102;&#101;&#103;&#97;&#110;&#64;&#108;&#108;&#114;&#46;&#105;&#110;&#50;&#112;&#51;&#46;&#102;&#114;
Laboratoire Leprince-Ringuet, CNRS/IN2P3, Ecole Polytechnique, Institut Polytechnique de Paris
This file is part of "calin". "calin" is free software: you can redistribute it and/or modify it under the
terms of the GNU General Public License version 2 or later, as published by
the Free Software Foundation. "calin" is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
Introduction
This notebook demonstrates how to generate PDF plots that sumamrize the diagnostics results written by compute_diagnostics.py.
First, import the required packages.
End of explanation
"""


sql = calin.io.sql_transceiver.SQLite3Transceiver("/CTA/diagnostics.sqlite",
  calin.io.sql_transceiver.SQLite3Transceiver.READ_ONLY)
diagnostics = calin.ix.scripts.compute_diagnostics.Results()
sql.retrieve_by_oid("diagnostics_results", 1, diagnostics)
del sql

"""
Explanation: Get diagnostics results from SQL database
Open the SQL diagnostics results previously written by compute_diagnostics.py and load the results. The diagnostics results are all assembled in a type calin.ix.scripts.compute_diagnostics.Results and written to a table diagnostics_results in the database. We thus:
1. Open the SQLite3 data file
2. Create a calin.ix.scripts.compute_diagnostics.Results instance to hold the results
3. Retrieve the results, assuming they are written in row 1 of the database
4. Close the database
End of explanation
"""


wfs  = diagnostics.waveform_stats()
psd  = diagnostics.waveform_psd()
sig  = diagnostics.sig_stats()
bkg  = diagnostics.bkg_stats()
smb  = diagnostics.sig_minus_bkg_stats()
t0   = diagnostics.t0_stats()
conf = diagnostics.run_config()
clo  = diagnostics.command_line_options()

"""
Explanation: Extract the elements that are of interest
See the file proto/scripts/compute_diagnostics.proto for description of the Results type.
End of explanation
"""


def plot_functional_stats(pdf, stats, stats_label, stats_name, stats_units):
    figure(figsize=(8.27, 11.69),dpi=300)
    gs = GridSpec(4, 2,
            height_ratios=[1,1,1.5,1.0],
            width_ratios=[1,1],
            hspace=0.4,
            wspace=0.3,
            bottom=0.05,
            top=0.9
          )

    figtext(0.5,0.97,'--- %s ---'%stats_label,va='top',ha='center',fontsize=16)
    figtext(0.3,0.94,'-- High gain -- ',va='top',ha='center',fontsize=14)
    figtext(0.725,0.94,'-- Low gain --',va='top',ha='center',fontsize=14)

    subplot(gs[0,0])
    title("Histogram of mean over all channels")
    ylabel('Count [events/%.2f %s]'%(stats.high_gain().mean_hist().dxval(),stats_units))
    xlabel('Mean %s'%stats_name)
    plot_histogram(stats.high_gain().mean_hist())

    subplot(gs[0,1])
    title("Histogram of mean over all channels")
    ylabel('Count [events/%.2f %s]'%(stats.low_gain().mean_hist().dxval(),stats_units))
    xlabel('Mean %s'%stats_name)
    plot_histogram(stats.low_gain().mean_hist())

    subplot(gs[1,0])
    title("Channel mean and RMS")
    ylabel('%s [%s]'%(stats_name,stats_units))
    xlabel('Channel index')
    xmean = calin.diagnostics.functional.channel_mean(stats.high_gain())
    xrms = sqrt(calin.diagnostics.functional.channel_var(stats.high_gain()))
    errorbar(frange(0,len(xmean),closed=False),xmean,xrms,fmt='.')

    subplot(gs[1,1])
    title("Channel mean and RMS")
    ylabel('%s [%s]'%(stats_name,stats_units))
    xlabel('Channel index')
    xmean = calin.diagnostics.functional.channel_mean(stats.low_gain())
    xrms = sqrt(calin.diagnostics.functional.channel_var(stats.low_gain()))
    errorbar(frange(0,len(xmean),closed=False),xmean,xrms,fmt='.')

    subplot(gs[2,0])
    h=pcolor(calin.diagnostics.functional.channel_cov_frac(stats.high_gain()),
            cmap=get_cmap('jet'));
    title("Channel correlation matrix")
    xlabel('Channel index')
    ylabel('Channel index')
    cb = colorbar()
    axis('square')

    subplot(gs[2,1])
    h=pcolor(calin.diagnostics.functional.channel_cov_frac(stats.low_gain()),
            cmap=get_cmap('jet'));
    title("Channel correlation matrix")
    xlabel('Channel index')
    ylabel('Channel index')
    cb = colorbar()
    axis('square')

    subplot(gs[3,:])
    title("High-gain to low-gain correlation")
    ylabel('Fraction [1]')
    xlabel('Channel index')
    xfrac = calin.diagnostics.functional.channel_high_to_low_gain_cov_frac(stats)
    plot(frange(0,len(xfrac),closed=False),xfrac,'.')
    
    pdf.savefig()
    close(gcf())

"""
Explanation: Plot the results as PDF pages
The following functions (which I don't describe in detail) plot various elements of the diagnostics onto pages of PDF document. They use various helper functions to convert some of the elements into means, variances and covariance matrices.
Plot functional statistics summary
This can be called to plot the statistics for any of the functional values in the diagnostics data, including pulse start time, integrated signal or background etc.
For the functional value calculated from the high and low gain waveforms plot:
1) Histogram of mean across all channels
2) Mean and RMS for each channel
3) Channel-to-channel correlation
And finish by plotting the correlation between the high and low-gain values for each channel.
End of explanation
"""


def plot_wf_summary_page(pdf):
#    colcycle = [ '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2' ]
    colcycle = [ '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#bcbd22', '#17becf' ]

    figure(figsize=(8.27, 11.69),dpi=300)
    gs = GridSpec(2, 2,
            height_ratios=[1,5],
            width_ratios=[1,1],
            hspace=0.2,
            wspace=0.3,
            bottom=0.05,
            top=0.9
          )

    figtext(0.5,0.97,'--- Waveforms ---',va='top',ha='center',fontsize=16)
    figtext(0.3,0.94,'-- High gain -- ',va='top',ha='center',fontsize=14)
    figtext(0.725,0.94,'-- Low gain --',va='top',ha='center',fontsize=14)
    
    lg_wf_mean = zeros(conf.num_samples())
    hg_wf_mean = zeros(conf.num_samples())
    for ich in  range(0,wfs.high_gain_size()):
        hg_wf_mean += calin.diagnostics.waveform.WaveformStatsVisitor.waveform_mean(wfs.high_gain(ich))
        lg_wf_mean += calin.diagnostics.waveform.WaveformStatsVisitor.waveform_mean(wfs.low_gain(ich))
    hg_wf_mean /= wfs.high_gain_size()
    lg_wf_mean /= wfs.high_gain_size()
        
    subplot(gs[0,0])
    title("All channel mean waveform")
    ylabel('Mean waveform [DC]')
    xlabel('Sample number')
    plot(hg_wf_mean,'k')
    a=axis()
    gca().add_patch(Rectangle((clo.sig_window_start()-0.5, a[2]), clo.window_size(), a[3]-a[2], facecolor='#ffeeee', edgecolor='none'))
    text(clo.sig_window_start()+(clo.window_size()-1)/2,a[2]*0.975+a[3]*0.025,'Sig',ha='center',va='bottom',color='r')
    gca().add_patch(Rectangle((clo.bkg_window_start()-0.5, a[2]), clo.window_size(), a[3]-a[2], facecolor='#eeeeff', edgecolor='none'))
    text(clo.bkg_window_start()+(clo.window_size()-1)/2,a[3]*0.975+a[2]*0.025,'Bkg',ha='center',va='top',color='b')

    subplot(gs[0,1])
    title("All channel mean waveform")
    ylabel('Mean waveform [DC]')
    xlabel('Sample number')
    plot(lg_wf_mean,'k')
    a=axis()
    gca().add_patch(Rectangle((clo.sig_window_start()-0.5, a[2]), clo.window_size(), a[3]-a[2], facecolor='#ffeeee', edgecolor='none'))
    text(clo.sig_window_start()+(clo.window_size()-1)/2,a[2]*0.975+a[3]*0.025,'Sig',ha='center',va='bottom',color='r')
    gca().add_patch(Rectangle((clo.bkg_window_start()-0.5, a[2]), clo.window_size(), a[3]-a[2], facecolor='#eeeeff', edgecolor='none'))
    text(clo.bkg_window_start()+(clo.window_size()-1)/2,a[3]*0.975+a[2]*0.025,'Bkg',ha='center',va='top',color='b')

    subplot(gs[1,0])
    gca().set_prop_cycle(cycler('color', colcycle) + cycler('lw', [1.5, 1, 1, 1, 1, 1, 1]))
    title("Mean waveforms")
    ylabel('Channel number + scaled waveform')
    xlabel('Sample index')
    max_y = -inf
    for ich in  range(0,wfs.high_gain_size()):
        wf = calin.diagnostics.waveform.WaveformStatsVisitor.waveform_mean(wfs.high_gain(ich))
        wf -= wf[0]
        wf = ich + wf/(max(hg_wf_mean)-min(hg_wf_mean))*3
        plot(wf,'-')
        max_y = max(max_y, max(wf))
        a=list(axis());
        a[2] = -0.5
        a[3] = max_y+0.5
        axis(a)

    subplot(gs[1,1])
    gca().set_prop_cycle(cycler('color', colcycle) + cycler('lw', [1.5, 1, 1, 1, 1, 1, 1]))
    title("Mean waveforms")
    ylabel('Channel number + scaled waveform')
    xlabel('Sample index')
    max_y = -inf
    for ich in  range(0,wfs.low_gain_size()):
        wf = calin.diagnostics.waveform.WaveformStatsVisitor.waveform_mean(wfs.low_gain(ich))
        wf -= wf[0]
        wf = ich + wf/(max(lg_wf_mean)-min(lg_wf_mean))*3
        plot(wf,'-')
        max_y = max(max_y, max(wf))
        a=list(axis());
        a[2] = -0.5
        a[3] = max_y+0.5
        axis(a)

    pdf.savefig()
    close(gcf())

"""
Explanation: Plot waveform summary

Mean trace over all channels
Multi-channel "oscilloscope-like" display of mean trace for all channels, seperated out with small gap
End of explanation
"""


def plot_channel_page(pdf, ich):
    hg_wfs = wfs.high_gain(ich)
    hg_psd = psd.high_gain(ich)
    hg_smb = smb.high_gain().value_hist(ich)
    hg_t0  = t0.high_gain().value_hist(ich)
    lg_wfs = wfs.low_gain(ich)
    lg_psd = psd.low_gain(ich)
    lg_smb = smb.low_gain().value_hist(ich)
    lg_t0  = t0.low_gain().value_hist(ich)

    chid = conf.configured_channel_id(ich)
    ch = conf.camera_layout().channel(chid)

    figure(figsize=(8.27, 11.69),dpi=300)
    gs = GridSpec(5, 2,
            height_ratios=[1,1,1,1.5,1],
            width_ratios=[1,1],
            hspace=0.4,
            wspace=0.3,
            bottom=0.05,
            top=0.9
          )

    figtext(0.5,0.97,'--- Channel %d (mod: %d/%d) ---'%(ch.channel_index(),
        ch.module_index(),ch.module_channel_index()),va='top',ha='center',fontsize=16)
    figtext(0.3,0.94,'-- High gain -- ',va='top',ha='center',fontsize=14)
    figtext(0.725,0.94,'-- Low gain --',va='top',ha='center',fontsize=14)

    subplot(gs[0,0])
    title("Mean waveform")
    ylabel('Mean waveform [DC]')
    xlabel('Sample number')
    plot(calin.diagnostics.waveform.WaveformStatsVisitor.waveform_mean(hg_wfs))
    a=axis()
    if(a[3] >= 4095):
        axhline(y=4095,linestyle='--',color='k')
    a=axis()
    gca().add_patch(Rectangle((clo.sig_window_start()-0.5, a[2]), clo.window_size(), a[3]-a[2], facecolor='#ffeeee', edgecolor='none'))
    text(clo.sig_window_start()+(clo.window_size()-1)/2,a[2]*0.975+a[3]*0.025,'Sig',ha='center',va='bottom',color='r')
    gca().add_patch(Rectangle((clo.bkg_window_start()-0.5, a[2]), clo.window_size(), a[3]-a[2], facecolor='#eeeeff', edgecolor='none'))
    text(clo.bkg_window_start()+(clo.window_size()-1)/2,a[3]*0.975+a[2]*0.025,'Bkg',ha='center',va='top',color='b')

    subplot(gs[0,1])
    title("Mean waveform")
    ylabel('Mean waveform [DC]')
    xlabel('Sample number')
    plot(calin.diagnostics.waveform.WaveformStatsVisitor.waveform_mean(lg_wfs))
    a=axis()
    if(a[3] >= 4095):
        axhline(y=4095,linestyle='--',color='k')
    a=axis()
    gca().add_patch(Rectangle((clo.sig_window_start()-0.5, a[2]), clo.window_size(), a[3]-a[2], facecolor='#ffeeee', edgecolor='none'))
    text(clo.sig_window_start()+(clo.window_size()-1)/2,a[2]*0.975+a[3]*0.025,'Sig',ha='center',va='bottom',color='r')
    gca().add_patch(Rectangle((clo.bkg_window_start()-0.5, a[2]), clo.window_size(), a[3]-a[2], facecolor='#eeeeff', edgecolor='none'))
    text(clo.bkg_window_start()+(clo.window_size()-1)/2,a[3]*0.975+a[2]*0.025,'Bkg',ha='center',va='top',color='b')

    subplot(gs[1,0])
    title("Signal minus background")
    ylabel('Frequency [events/DC]')
    xlabel('Integrated charge [DC samples]')
    plot_histogram(hg_smb)

    subplot(gs[1,1])
    title("Signal minus background")
    ylabel('Frequency [events/DC]')
    xlabel('Integrated charge [DC samples]')
    plot_histogram(lg_smb)

    subplot(gs[2,0])
    title("Pulse start time (T0)")
    ylabel('Frequency [events/DC]')
    xlabel('Pulse start time [samples]')
    plot_histogram(hg_t0)

    subplot(gs[2,1])
    title("Pulse start time (T0)")
    ylabel('Frequency [events/DC]')
    xlabel('Pulse start time [samples]')
    plot_histogram(lg_t0)

    subplot(gs[3,0])
    h=pcolor(calin.diagnostics.waveform.WaveformStatsVisitor.waveform_cov_frac(hg_wfs),
            cmap=get_cmap('bwr'));
    title("Waveform covariance matrix")
    xlabel('Sample number')
    ylabel('Sample number')
    h.set_clim((-1.0,1.0))
    cb = colorbar()
    cb.set_label("Covariance fraction")
    axis('square')

    subplot(gs[3,1])
    h=pcolor(calin.diagnostics.waveform.WaveformStatsVisitor.waveform_cov_frac(lg_wfs),
            cmap=get_cmap('bwr'));
    title("Waveform covariance matrix")
    xlabel('Sample number')
    ylabel('Sample number')
    h.set_clim((-1.0,1.0))
    cb = colorbar()
    cb.set_label("Covariance fraction [1]")
    axis('square')

    subplot(gs[4,0])
    f = asarray(range(0,hg_psd.psd_sum_size()))/60*1000
    semilogy(f, calin.diagnostics.waveform.WaveformPSDVisitor.psd_mean(hg_psd))
    title("High-gain PSD")
    ylabel('Signal squared [DC^2]')
    xlabel('Frequency [MHz]')

    subplot(gs[4,1])
    c, x = calin.diagnostics.waveform.WaveformPSDVisitor.corr_mean_centered(hg_psd, hg_wfs)
    plot(x,c)
    title("High-gain auto correlation")
#        ylabel('Signal squared [DC^2]')
    xlabel('Time shift [samples]')

    pdf.savefig()
    close(gcf())

"""
Explanation: Plot per-channel statistics page
For the high-gain and low-gain waveforms plot the:

Mean waveform across all events
Histogram of signal minus background for the channel
Histogram of the pulse start time for the channel
Waveform sample-to-sample covariance matrix for the channel

Finally, also plot the high-gain mean power-spectral density and auto-correlation.
End of explanation
"""


with PdfPages('/CTA/diagnostics.pdf') as pdf:
    plot_wf_summary_page(pdf)
    plot_functional_stats(pdf, t0, 'Pulse start time', 't0', 'samples')
    plot_functional_stats(pdf, sig, 'Signal integration', 'Charge', 'DC')
    plot_functional_stats(pdf, bkg, 'Background integration', 'Charge', 'DC')
    plot_functional_stats(pdf, smb, 'Signal minus background', 'Charge', 'DC')

    for ich in range(0,wfs.high_gain_size()):
        print(ich)
        plot_channel_page(pdf, ich)
            
    d = pdf.infodict()
    d['Title'] = 'Diagnostics'
    d['CreationDate'] = datetime.datetime.today()

"""
Explanation: Generate the pages of the PDF file

The waveform summary
Statistics of the pulse start time (t0)
Statistics of the signal integration window
Statistics of the background integration window
Statistics of the signal minus the background integration window
One page for each channel
End of explanation
"""


