content: #!/usr/bin/env python
  """
  Do the Single Photoelectron anaylsis
  
  Usage:
    digicam-spe [options] [--] <INPUT>...
  
  Options:
    -h --help                    Show this screen.
    --max_events=N               Maximum number of events to analyse.
    --max_histo_filename=FILE    File path of the max histogram.
                                 [Default: ./max_histo.pk]
    --charge_histo_filename=FILE File path of the charge histogram
                                 [Default: ./charge_histo.pk]
    --raw_histo_filename=FILE    File path of the raw histogram
                                 [Default: ./raw_histo.pk]
    -o OUTPUT --output=OUTPUT    Output file path to store the results.
                                 [Default: ./results.npz]
    -c --compute                 Compute the data.
    -f --fit                     Fit.
    -d --display                 Display.
    -v --debug                   Enter the debug mode.
    -p --pixel=<PIXEL>           Give a list of pixel IDs.
    --shift=N                    Number of bins to shift before integrating
                                 [default: 0].
    --integral_width=N           Number of bins to integrate over
                                 [default: 7].
    --pulse_finder_threshold=F   Threshold of pulse finder in arbitrary units
                                 [default: 2.0].
    --save_figures=PATH          Save the plots to the indicated folder.
                                 Figures are not saved is set to none
                                 [default: none]
    --ncall=N                    Number of calls for the fit [default: 10000]
    --n_samples=N                Number of samples per waveform
  
  """
  import os
  
  import matplotlib.pyplot as plt
  import numpy as np
  from docopt import docopt
  from histogram.histogram import Histogram1D
  from tqdm import tqdm
  
  from digicampipe.calib.baseline import fill_baseline, subtract_baseline
  from digicampipe.calib.charge import compute_charge
  from digicampipe.calib.peak import find_pulse_with_max, \
      find_pulse_fast
  from digicampipe.io.event_stream import calibration_event_stream
  from digicampipe.scripts import raw
  from digicampipe.scripts.fmpe import FMPEFitter
  from digicampipe.utils.docopt import convert_pixel_args, \
      convert_int, convert_text
  from digicampipe.utils.pdf import fmpe_pdf_10
  
  
  class MaxHistoFitter(FMPEFitter):
      def __init__(self, histogram, estimated_gain, **kwargs):
          n_peaks = 2
          super(MaxHistoFitter, self).__init__(histogram, estimated_gain,
                                               n_peaks, **kwargs)
          self.parameters_plot_name = {'baseline': '$B$', 'gain': 'G',
                                       'sigma_e': '$\sigma_e$',
                                       'sigma_s': '$\sigma_s$',
                                       'a_0': None, 'a_1': None}
  
      def pdf(self, x, baseline, gain, sigma_e, sigma_s, a_0, a_1):
          params = {'baseline': baseline, 'gain': gain, 'sigma_e': sigma_e,
                    'sigma_s': sigma_s, 'a_0': a_0, 'a_1': a_1, 'bin_width': 0}
  
          return fmpe_pdf_10(x, **params)
  
  
  class SPEFitter(FMPEFitter):
      def __init__(self, histogram, estimated_gain, **kwargs):
          n_peaks = 4
          super(SPEFitter, self).__init__(histogram, estimated_gain, n_peaks,
                                          **kwargs)
          self.parameters_plot_name = {'baseline': '$B$', 'gain': 'G',
                                       'sigma_e': '$\sigma_e$',
                                       'sigma_s': '$\sigma_s$',
                                       'a_1': None, 'a_2': None, 'a_3': None,
                                       'a_4': None}
  
      def pdf(self, x, baseline, gain, sigma_e, sigma_s, a_1, a_2, a_3, a_4):
          params = {'baseline': baseline, 'gain': gain, 'sigma_e': sigma_e,
                    'sigma_s': sigma_s, 'a_0': 0, 'a_1': a_1, 'a_2': a_2,
                    'a_3': a_3, 'a_4': a_4, 'bin_width': 0}
  
          return fmpe_pdf_10(x, **params)
  
      def initialize_fit(self):
          init_params = super(SPEFitter, self).initialize_fit()
  
          init_params['a_4'] = init_params['a_3']
          init_params['a_3'] = init_params['a_2']
          init_params['a_2'] = init_params['a_1']
          init_params['a_1'] = init_params['a_0']
  
          init_params['baseline'] = init_params['baseline'] - init_params['gain']
  
          del init_params['a_0']
  
          self.initial_parameters = init_params
  
          return init_params
  
  
  def compute_dark_rate(number_of_zeros, total_number_of_events, time):
      p_0 = number_of_zeros / total_number_of_events
      rate = - np.log(p_0)
      rate /= time
  
      return rate
  
  
  def compute_max_histo(files, histo_filename, pixel_id, max_events,
                        integral_width, shift, baseline):
      n_pixels = len(pixel_id)
  
      if not os.path.exists(histo_filename):
  
          events = calibration_event_stream(files, pixel_id=pixel_id,
                                            max_events=max_events)
          # events = compute_baseline_with_min(events)
          events = fill_baseline(events, baseline)
          events = subtract_baseline(events)
          events = find_pulse_with_max(events)
          events = compute_charge(events, integral_width, shift)
          max_histo = Histogram1D(
              data_shape=(n_pixels,),
              bin_edges=np.arange(-4095 * integral_width,
                                  4095 * integral_width),
          )
  
          for event in events:
              max_histo.fill(event.data.reconstructed_charge)
  
          max_histo.save(histo_filename)
  
          return max_histo
  
      else:
  
          max_histo = Histogram1D.load(histo_filename)
  
          return max_histo
  
  
  def compute_spe(files, histo_filename, pixel_id, baseline, max_events,
                  integral_width, shift, pulse_finder_threshold, debug=False):
      if not os.path.exists(histo_filename):
  
          n_pixels = len(pixel_id)
          events = calibration_event_stream(files,
                                            max_events=max_events,
                                            pixel_id=pixel_id)
  
          events = fill_baseline(events, baseline)
          events = subtract_baseline(events)
          # events = find_pulse_1(events, 0.5, 20)
          # events = find_pulse_2(events, widths=[5, 6], threshold_sigma=2)
          events = find_pulse_fast(events, threshold=pulse_finder_threshold)
          # events = find_pulse_fast_2(events, threshold=pulse_finder_threshold,
          #                           min_dist=3)
          # events = find_pulse_correlate(events,
          #                               threshold=pulse_finder_threshold)
          # events = find_pulse_gaussian_filter(events,
          #                                    threshold=pulse_finder_threshold)
  
          # events = find_pulse_wavelets(events, widths=[4, 5, 6],
          #                             threshold_sigma=2)
  
          events = compute_charge(events, integral_width=integral_width,
                                  shift=shift)
          # events = compute_amplitude(events)
          # events = fit_template(events)
  
          # events = compute_full_waveform_charge(events)
  
          spe_histo = Histogram1D(
              data_shape=(n_pixels,),
              bin_edges=np.arange(-4095 * 50, 4095 * 50)
          )
  
          for event in events:
              spe_histo.fill(event.data.reconstructed_charge)
  
          spe_histo.save(histo_filename)
  
          return spe_histo
  
      else:
  
          spe_histo = Histogram1D.load(histo_filename)
  
          return spe_histo
  
  
  def entry():
      args = docopt(__doc__)
  
      files = args['<INPUT>']
      debug = args['--debug']
  
      max_events = convert_int(args['--max_events'])
  
      raw_histo_filename = args['--raw_histo_filename']
      charge_histo_filename = args['--charge_histo_filename']
      max_histo_filename = args['--max_histo_filename']
      results_filename = args['--output']
  
      pixel_id = convert_pixel_args(args['--pixel'])
      n_pixels = len(pixel_id)
  
      integral_width = int(args['--integral_width'])
      shift = int(args['--shift'])
      pulse_finder_threshold = float(args['--pulse_finder_threshold'])
  
      n_samples = int(args['--n_samples'])  # TODO access this in a better way !
      estimated_gain = 20
      ncall = int(args['--ncall'])
  
      if args['--compute']:
          raw_histo = raw.compute(files, max_events=max_events,
                                  pixel_id=pixel_id, filename=raw_histo_filename)
          baseline = raw_histo.mode()
  
          compute_max_histo(files, max_histo_filename, pixel_id, max_events,
                            integral_width, shift, baseline)
  
          compute_spe(files, charge_histo_filename, pixel_id, baseline,
                      max_events, integral_width, shift, pulse_finder_threshold,
                      debug=debug)
  
      if args['--fit']:
  
          spe_histo = Histogram1D.load(charge_histo_filename)
          max_histo = Histogram1D.load(max_histo_filename)
  
          dark_count_rate = np.zeros(n_pixels) * np.nan
          electronic_noise = np.zeros(n_pixels) * np.nan
          crosstalk = np.zeros(n_pixels) * np.nan
          gain = np.zeros(n_pixels) * np.nan
  
          for i, pixel in tqdm(enumerate(pixel_id), total=n_pixels,
                               desc='Pixel'):
              histo = max_histo[i]
              fitter = MaxHistoFitter(histo, estimated_gain, throw_nan=True)
  
              try:
                  fitter.fit(ncall=100)
                  fitter.fit(ncall=ncall)
  
                  n_entries = histo.data.sum()
                  number_of_zeros = fitter.parameters['a_0']
                  window_length = 4 * n_samples
                  rate = compute_dark_rate(number_of_zeros,
                                           n_entries,
                                           window_length)
                  electronic_noise[i] = fitter.parameters['sigma_e']
  
                  dark_count_rate[i] = rate
                  if debug:
                      fitter.draw()
                      fitter.draw_init(x_label='[LSB]')
                      fitter.draw_fit(x_label='[LSB]')
                      plt.show()
  
              except Exception as e:
  
                  print('Could not compute dark count rate'
                        ' in pixel {}'.format(pixel))
                  print(e)
  
          np.savez(results_filename, dcr=dark_count_rate,
                   sigma_e=electronic_noise, pixel_id=pixel_id)
  
          for i, pixel in tqdm(enumerate(pixel_id), total=n_pixels,
                               desc='Pixel'):
  
              histo = spe_histo[i]
              fitter = SPEFitter(histo, estimated_gain, throw_nan=True)
  
              try:
  
                  fitter.fit(ncall=100)
                  fitter.fit(ncall=ncall)
  
                  params = fitter.parameters
                  n_entries = params['a_1']
                  n_entries += params['a_2']
                  n_entries += params['a_3']
                  n_entries += params['a_4']
                  crosstalk[i] = (n_entries - params['a_1']) / n_entries
                  gain[i] = params['gain']
  
                  if debug:
                      fitter.draw()
                      fitter.draw_init(x_label='[LSB]')
                      fitter.draw_fit(x_label='[LSB]')
                      plt.show()
  
              except Exception as e:
  
                  print('Could not compute gain and crosstalk'
                        ' in pixel {}'.format(pixel))
                  print(e)
  
          data = dict(np.load(results_filename))
          data['crosstalk'] = crosstalk
          data['gain'] = gain
          np.savez(results_filename, **data)
  
      save_figure = convert_text(args['--save_figures'])
      if save_figure is not None:
          output_path = save_figure
          spe_histo = Histogram1D.load(charge_histo_filename)
          spe_amplitude = Histogram1D.load(charge_histo_filename)
          raw_histo = Histogram1D.load(raw_histo_filename)
          max_histo = Histogram1D.load(max_histo_filename)
  
          figure_directory = output_path + 'figures/'
  
          if not os.path.exists(figure_directory):
              os.makedirs(figure_directory)
  
          histograms = [spe_histo, spe_amplitude, raw_histo, max_histo]
          names = ['histogram_charge/', 'histogram_amplitude/', 'histogram_raw/',
                   'histo_max/']
  
          for i, histo in enumerate(histograms):
  
              figure = plt.figure()
              histogram_figure_directory = figure_directory + names[i]
  
              if not os.path.exists(histogram_figure_directory):
                  os.makedirs(histogram_figure_directory)
  
              for j, pixel in enumerate(pixel_id):
                  axis = figure.add_subplot(111)
                  figure_path = histogram_figure_directory + 'pixel_{}'. \
                      format(pixel)
  
                  try:
  
                      histo.draw(index=(j,), axis=axis, log=True, legend=False)
                      figure.savefig(figure_path)
  
                  except Exception as e:
  
                      print('Could not save pixel {} to : {} \n'.
                            format(pixel, figure_path))
                      print(e)
  
                  axis.remove()
  
      if args['--display']:
  
          spe_histo = Histogram1D.load(charge_histo_filename)
          raw_histo = Histogram1D.load(os.path.join(output_path,
                                                    raw_histo_filename))
          max_histo = Histogram1D.load(max_histo_filename)
  
          spe_histo.draw(index=(0,), log=True, legend=False)
          raw_histo.draw(index=(0,), log=True, legend=False)
          max_histo.draw(index=(0,), log=True, legend=False)
  
          try:
  
              data = np.load(results_filename)
              dark_count_rate = data['dcr']
              electronic_noise = data['sigma_e']
              crosstalk = data['crosstalk']
              gain = data['gain']
  
          except IOError as e:
  
              print(e)
              print('Could not find the analysis files !')
  
          plt.figure()
          plt.hist(dark_count_rate[np.isfinite(dark_count_rate)],
                   bins='auto')
          plt.xlabel('dark count rate [GHz]')
          plt.legend(loc='best')
  
          plt.figure()
          plt.hist(crosstalk[np.isfinite(crosstalk)],
                   bins='auto')
          plt.xlabel('Crosstalk []')
          plt.legend(loc='best')
  
          plt.figure()
          plt.hist(gain[np.isfinite(gain)],
                   bins='auto')
          plt.xlabel('Gain [LSB/p.e.]')
          plt.legend(loc='best')
  
          plt.figure()
          plt.hist(electronic_noise[np.isfinite(electronic_noise)],
                   bins='auto')
          plt.xlabel('$\sigma_e$ [LSB]')
          plt.legend(loc='best')
  
          plt.show()
  
      return
  
  
  if __name__ == '__main__':
      entry()
  