# -*- coding: utf-8 -*-
"""
Created on Mon Jun  8 15:37:51 2015

@author: Anton O Lindahl
"""
import h5py
import argparse
import matplotlib.pyplot as plt
import numpy as np
import time
import os
import sys
import lmfit
import warnings

from aolPyModules import wiener, wavelet_filter
import time_to_energy_conversion as tof_to_energy
from aolPyModules import plotting as aol_plotting
import area_fill

prompt_roi = [1.508, 1.535]
streak_time_roi = [1.57, 1.66]
wt_th = 0.03
energy_scale_eV = np.linspace(40, 160, 2**9)
time_stamp = 'time_stamp'
data_dir = 'h5_files'
h5_file_name_template = data_dir + '/run{}_all.h5'
response_file_name = data_dir + '/response.h5'
nois_file_name = data_dir + '/noise.h5'
tof_to_energy_conversion_file_name = data_dir + '/time_to_energy.h5'


def h5_file_name_funk(run):
    return h5_file_name_template.format(run)


def update_progress(i_evt, n_events, verbose=True):
    if (verbose and
            ((i_evt % (n_events / 100) == 0) or (i_evt == n_events-1))):
        progress = (100 * i_evt) / (n_events - 1)
        num_squares = 40
        base_string = '\r[{:' + str(num_squares) + '}] {}%'
        print base_string.format('#'*(progress * num_squares / 100), progress),
        sys.stdout.flush()


def list_hdf5_content(group, indent='  '):
    for k, v in group.iteritems():
        print '{}"{}"'.format(indent, k),
        if isinstance(v, h5py.Group):
            print 'group with members:'
            list_hdf5_content(v, indent=indent + '  ')
        elif isinstance(v, h5py.Dataset):
            print '\t{} {}'.format(v.shape, v.dtype)


def make_dataset(h5, name, shape, dtype=np.float):
    try:
        dset = h5.require_dataset(name, shape=shape,
                                  dtype=dtype, exact=True)
    except TypeError:
        del h5[name]
        dset = h5.create_dataset(name, shape=shape, dtype=np.float)
    if time_stamp not in dset.attrs.keys():
        dset.attrs.create(time_stamp, 0)
    return dset


def make_group(h5, name):
    try:
        group = h5.require_group(name)
    except TypeError:
        del h5[name]
        group = h5.create_group(name)
    if time_stamp not in group.attrs.keys():
        group.attrs.create(time_stamp, 0)
    return group


def older(dset, dset_list):
    if (isinstance(dset_list, h5py.Dataset) or
            isinstance(dset_list, h5py.Group)):
        return dset.attrs[time_stamp] < dset_list.attrs[time_stamp]

    return np.any([dset.attrs[time_stamp] < d.attrs[time_stamp] for
                   d in dset_list])


class Timer_object:
    def __init__(self, t):
        self.attrs = {'time_stamp': t}


class Tims_stamp_warning(Warning):
    pass


def time_stamp_object(h5_object):
    try:
        h5_object.attrs['time_stamp'] = time.time()
    except:
        warnings.warn('Could not time stamp the object {}.'.format(
            repr(h5_object)))


def get_response(plot=False, verbose=0):
    try:
        with h5py.File(response_file_name, 'r') as f:
            response = f['signal'].value
            t = f['signal'].attrs[time_stamp]
    except IOError:
        if verbose > 0:
            print 'Could not open response file. Trying to make it.'
        response, t = construct_response(verbose=verbose)
    if plot:
        with h5py.File(response_file_name, 'r') as f:
            time_scale = f['time_scale'].value
        plt.figure('response')
        plt.clf()
        plt.plot(time_scale, response)

    return response, t


def construct_response(plot=False, verbose=0):
    # The Kr runs
    runs = [132, 133, 134, 135, 136]
    if verbose > 0:
        print 'Loading Kr files for prompt determination.'
    h5_file_names = [h5_file_name_template.format(run) for run in runs]
    h5_list = []
    for file_name in h5_file_names:
        update_run_contained_derived_data(file_name, verbose=verbose)
        h5_list.append(h5py.File(file_name, 'r+'))
    time_scale = h5_list[0]['raw/time_scale'].value
    response = np.zeros_like(time_scale)
    n_shots = 0
    sl = slice(time_scale.searchsorted(prompt_roi[0]),
               time_scale.searchsorted(prompt_roi[1], side='right'))
    for h5 in h5_list:
        response[sl] += h5['raw/time_signal'][:, sl].sum(0)
        n_shots += h5['raw/event_time_s'].shape[0]
    response /= n_shots

    response[sl] = wiener.edgeSmoothing(response[sl], smoothPoints=15)

    response /= response.sum()

    with h5py.File(response_file_name, 'w') as res_file:
        dset = res_file.create_dataset('signal', data=response)
        dset.attrs.create(time_stamp, time.time())
        res_file.create_dataset('time_scale', data=time_scale)

    return get_response(plot=plot, verbose=verbose)


def get_file_names_for_noise_spectrum():
    return ['/'.join([data_dir, f]) for f in os.listdir(data_dir) if
            f.startswith('run') and f.endswith('_all.h5')]


def get_nois_spectrum(plot=False, verbose=0):
    try:
        with h5py.File(nois_file_name, 'r') as f:
            pass
        new_noise = False
    except IOError:
        if verbose > 0:
            print 'Could not open response file. Trying to make it.',
            print 'In "get_nois_spectrum()".'
        construct_nois_spectrum(plot=plot, verbose=verbose)
        new_noise = True

    if not new_noise:
        make_new_noise = False
        with h5py.File(nois_file_name, 'r') as f:
            noise = f['noise']
            h5_file_names = get_file_names_for_noise_spectrum()
            for h5_name in h5_file_names:
                with h5py.File(h5_name, 'r') as h5:
                    if older(noise, h5['raw']):
                        make_new_noise = True
                        if verbose > 0:
                            print 'Noise was made earlier than the raw data',
                            print 'in the file', h5_name, 'Make new noise.'
                        break
                    elif False:
                        print 'Noise was made later than the raw data in',
                        print 'the file', h5_name

        if make_new_noise:
            construct_nois_spectrum(plot=plot, verbose=verbose)

    with h5py.File(nois_file_name, 'r') as f:
        noise = f['noise']
        return noise.value, noise.attrs['time_stamp']


def construct_nois_spectrum(plot=False, verbose=0):
    h5_file_names = get_file_names_for_noise_spectrum()
    for file_name in h5_file_names:
        update_run_contained_derived_data(file_name)
    empty_shots = []
    for i, h5_name in enumerate(h5_file_names):
        with h5py.File(h5_name, 'r') as h5:
            time_signal_dset = h5['raw/time_signal']
            try:
                max_signal = h5['max_signal'].value
            except KeyError:
                max_signal = np.max(time_signal_dset.value, axis=1)
            no_x_rays = max_signal < 0.04
            if no_x_rays.sum() > 0:
                empty_shots.extend(time_signal_dset[no_x_rays, :])
            if i == 0:
                time_scale = h5['raw/time_scale'].value
            if verbose > 0:
                print h5_name, 'has', no_x_rays.sum(), 'empty shots'
    empty_shots = np.array(empty_shots)

#    print len(empty_shots)
#    plt.figure('snr')
#    plt.clf()
#    for shot in empty_shots[:]:
#        plt.plot(time_scale, shot)

    freq = (np.linspace(0., 1., len(time_scale)) *
            1e-3/(time_scale[1] - time_scale[0]))
    fft_empty_shots = np.fft.fft(empty_shots, axis=1)
    amp = np.mean(np.abs(fft_empty_shots)**2, axis=0)
    wt_amp = amp[:]
    wt_amp = wavelet_filter.wavelet_filt(amp[1:], thresh=wt_th)
    wt_amp[1:] = (wt_amp[1:] + wt_amp[-1:0:-1]) / 2

#    plt.figure('fft')
#    plt.clf()
#    plt.plot(freq, amp)
#    plt.plot(freq, wt_amp, 'r')

    with h5py.File(nois_file_name, 'w') as f:
        dset = f.create_dataset('noise', data=wt_amp)
        dset.attrs.create('time_stamp', time.time())
        f.create_dataset('freq', data=freq)

    return get_nois_spectrum()


def construct_snr_spectrum(h5, plot=False):
    noise, t = get_nois_spectrum()

    sig_spec = h5['fft_spectrum_mean'].value
    freq = h5['fft_freq_axis'].value

    wt_spec = wavelet_filter.wavelet_filt(sig_spec, thresh=wt_th)
    wt_spec[1:] = (wt_spec[1:] + wt_spec[-1:0:-1]) / 2

    snr = (wt_spec - noise) / noise

    if plot:
        plt.figure('signal and noise')
        plt.clf()
        plt.semilogy(freq, sig_spec, label='signal')
        plt.semilogy(freq, noise, label='noise')
        plt.semilogy(freq, wt_spec, label='wt signal')
        plt.semilogy(freq, snr, label='snr')
        plt.legend(loc='best')

    return snr


def check_tof_to_energy_conversion_matrix(plot=False, verbose=0):
    try:
        with h5py.File(tof_to_energy_conversion_file_name, 'r'):
            pass
    except IOError:
        if verbose > 0:
            print 'Could not open the file. Making the conversion matrix.'
        construc_tof_to_energy_conversion_matrix(plot=plot, verbose=verbose)

    _, h5_dict, _ = tof_to_energy.load_tof_to_energy_data(verbose=verbose)
    with h5py.File(tof_to_energy_conversion_file_name, 'r') as trans_h5:
        if not older(
                trans_h5['matrix'],
                [h5['streak_peak_integral'] for h5 in h5_dict.itervalues()] +
                [Timer_object(1437117486)]):
            return
    if verbose > 0:
        print 'Conversion to old, remaking it.'
    construc_tof_to_energy_conversion_matrix(plot=plot, verbose=verbose)


def construc_tof_to_energy_conversion_matrix(plot=False, verbose=0):
    M, t, E, time_to_energy_params, tof_prediction_params = \
        tof_to_energy.make_tof_to_energy_matrix(
            energy_scale_eV=energy_scale_eV, plot=plot, verbose=verbose)
    with h5py.File(tof_to_energy_conversion_file_name, 'w') as h5:
        dset = h5.create_dataset('matrix', data=M)
        dset.attrs.create('time_stamp', time.time())
        dset = h5.create_dataset('time_scale', data=t)
        dset.attrs.create('time_stamp', time.time())
        dset = h5.create_dataset('energy_scale_eV', data=E)
        dset.attrs.create('time_stamp', time.time())
        for k in time_to_energy_params:
            dset = h5.create_dataset(k, data=time_to_energy_params[k].value)
            dset.attrs.create('time_stamp', time.time())
        for k in tof_prediction_params:
            dset = h5.require_dataset(k, (), np.float)
            dset[()] = tof_prediction_params[k].value
            dset.attrs.create('time_stamp', time.time())


def open_hdf5_file(file_name, plot=False, verbose=0):
    try:
        # Open the file
        h5 = h5py.File(file_name, 'r+')
    except BaseException as e:
        print 'Could not open the specified hdf5 file "{}".'.format(
            file_name)
        print 'Message was: {}'.format(e.message)
        return -1
    return h5


def get_com(x, y):
    idx_l, idx_h = fwxm(x, y, 0.0, return_data='idx')
    sl = slice(idx_l, idx_h)
    return ((x[sl] * y[sl]).sum()) / (y[sl].sum())


def fwxm(x, y, fraction=0.5, return_data=''):
    y_max = y.max()
    idx_max = y.argmax()

    y_f = y_max * fraction

    for i in range(idx_max, -1, -1):
        if y[i] < y_f:
            idx_low = i
            break
    else:
        idx_low = idx_max

    for i in range(idx_max, len(x)):
        if y[i] < y_f:
            idx_high = i
            break
    else:
        idx_high = idx_max

    if return_data == 'idx':
        return idx_low, idx_high

    if return_data == 'limits':
        return x[idx_low], x[idx_high]

    return (x[idx_low] + x[idx_high]) / 2, x[idx_high] - x[idx_low]


def get_trace_bounds(x, y,
                     threshold=0.0, min_width=2,
                     energy_offset=0,
                     useRel=False, threshold_rel=0.5,
                     roi=slice(None)):

    amp = y[roi]
    scale = x[roi]
    dx = np.mean(np.diff(x))

    if useRel:
        threshold_temp = threshold_rel * np.max(amp[np.isfinite(amp)])
        if threshold_temp < threshold:
            return [np.nan] * 3
        else:
            threshold_V = threshold_temp
    else:
        threshold_V = threshold

    nPoints = np.round(min_width/dx)

    i_min = 0
    for i in range(1, amp.size):
        if amp[i] < threshold_V:
            i_min = i
            continue
        if i-i_min >= nPoints:
            break
    else:
        return [np.nan] * 3

    i_max = amp.size - 1
    for i in range(amp.size-1, -1, -1):
        if amp[i] < threshold_V:
            i_max = i
            continue
        if i_max-i >= nPoints:
            break
    else:
        return [np.nan] * 3

    if i_min == 0 and i_max == amp.size - 1:
        return [np.nan] * 3

    # print 'min =', min, 'max =', max
    val_max = (scale[i_max] + (threshold_V - amp[i_max]) *
               (scale[i_max] - scale[i_max - 1]) /
               (amp[i_max] - amp[i_max - 1]))
    val_min = (scale[i_min] + (threshold_V - amp[i_min]) *
               (scale[i_min + 1] - scale[i_min]) /
               (amp[i_min + 1] - amp[i_min]))

    return val_min, val_max, threshold_V


def update_run_contained_derived_data(file_name, plot=False, verbose=0):
    """Update derived data based on information only in given file.

    Add some derived datasetd to the hdf5 file based on the raw data in the
    file. The added datasets are:

    - Mean of the FEE gas detectors for each shot: fee_mean
    - Maximum TOF waveform signal for each shot: max_signal
    - Frequency spectrum averaged over all shots: fft_spectrum_mean
    - The corresponding frequency axis: fft_freq_axis
    - BC2 energy calculated from the beam position: energy_BC2_MeV
    - L3 energy corrected based on the BC2 energy: energy_L3_corrected_MeV
    """
    if verbose > 0:
        print 'Entering "update_run_contained_derived_data()" ',
        print 'with file_name={}'.format(file_name)

    h5 = open_hdf5_file(file_name, plot, verbose)
    raw_group = h5['raw']
    n_events = raw_group['event_time_s'].shape[0]

    # Make the fee data set
    raw_fee_dset = raw_group['FEE_energy_mJ']
    fee_mean_dset = make_dataset(h5, 'fee_mean', (n_events,))
    if older(fee_mean_dset, raw_group):
        if verbose > 0:
            print 'Updating fee mean dataset'
        fee_mean_dset[:] = raw_fee_dset[:, 0: 4].mean(1)
        fee_mean_dset.attrs[time_stamp] = time.time()

    # Make max signal dataset
    time_signal_dset = raw_group['time_signal']
    max_sig_dset = make_dataset(h5, 'max_signal', (n_events,))
    if older(max_sig_dset, raw_group):
        if verbose > 0:
            print 'Get the maximum signal for each shot.'
        max_sig_dset[:] = np.max(time_signal_dset, axis=1)
        max_sig_dset.attrs['time_stamp'] = time.time()

    # Make the frequency spectrum
    time_scale = raw_group['time_scale'].value
    spectrum_dset = make_dataset(h5, 'fft_spectrum_mean', time_scale.shape)
    if older(spectrum_dset, [raw_group, max_sig_dset]):
        if verbose > 0:
            print 'Compute the frequency spectrum of the data.'
        max_signal = max_sig_dset.value
        use = max_signal > np.sort(max_signal)[-500:][0]
        signal = time_signal_dset[use, :]
        spectrum_dset[:] = np.mean(np.abs(np.fft.fft(signal, axis=1))**2,
                                   axis=0)
        spectrum_dset.attrs['time_stamp'] = time.time()
    freq_axis_dset = make_dataset(h5, 'fft_freq_axis', time_scale.shape)
    if older(freq_axis_dset, raw_group):
        if verbose > 0:
            print 'Updating the frequency axis.'
        freq_axis_dset[:] = (np.linspace(0., 1e-3, len(time_scale)) /
                             (time_scale[1] - time_scale[0]))
        freq_axis_dset.attrs['time_stamp'] = time.time()

    # Calculate the BC2 energy
    bc2_energy_dset = make_dataset(h5, 'energy_BC2_MeV', (n_events, ))
    if older(bc2_energy_dset, raw_group):
        if verbose > 0:
            print 'Calculating BC2 energy for the bpm reading.'
        # Values comes from a mail from Timothy Maxwell
        # The nominal BC2 energy is 5 GeV (was at least when this data was
        # recorded). The measurement is the relative offset of the beam
        # position in a BPM. The dispersion value is -364.7 mm.
        bc2_energy_dset[:] = 5e3 * (1. - raw_group['position_BC2_mm'][:] /
                                    364.7)
        bc2_energy_dset.attrs['time_stamp'] = time.time()

    # Calculate the corrected L3 energy
    l3_energy_cor_dset = make_dataset(h5, 'energy_L3_corrected_MeV',
                                      (n_events, ))
    if older(l3_energy_cor_dset, [raw_group, bc2_energy_dset,
                                  Timer_object(1434096408)]):
        if verbose > 0:
            print 'Calculating corrected L3 energy.'
        l3_energy_cor_dset[:] = (raw_group['energy_L3_MeV'][:] -
                                 (bc2_energy_dset[:] - 5000))
        l3_energy_cor_dset.attrs['time_stamp'] = time.time()

    # Make the phase cavity time filter
    pct_filter_dset = make_dataset(h5, 'pct_filter', (n_events, ),
                                   dtype=bool)
    if older(pct_filter_dset, [raw_group, Timer_object(0)]):
        print h5.filename
        pct0 = raw_group['phase_cavity_times'][:, 0]
        pct_filter_dset[:] = (0.4 < pct0) & (pct0 < 1.2)
        pct_filter_dset.attrs[time_stamp] = time.time()

    h5.close()


def update_with_noise_and_response(file_name, plot=False, verbose=0):
    """Update derived data based on noise and response spectra.

    Noise spectrum and detector response are determined form many runs. With
    these spectra a number of new paramters can be derived. These are:

    - snr_spectrum: Signal to Noise ratio spectrum based on the given noise \
    spectrum and the average spectrum in the current run.
    - filtered_time_signal: Wiegner deconvolution of the time signal based on \
    the signal to noise ratio and the detector response function.
    - streak_peak_center: Center of the streaking peak in the sense of the \
    center of mass of the peak in a given ROI. Based on the deconvoluted \
    signal.
    - streak_peak_integral: Photoline intensity by integration of the \
    deconvoluted spectrum in time domain.
    """

    # Make sure that the run contained information is up to date.
    update_run_contained_derived_data(file_name, plot, verbose-1)

    # Open the file.
    h5 = open_hdf5_file(file_name, plot, verbose)
    raw_group = h5['raw']
    n_events = raw_group['event_time_s'].shape[0]
    time_scale = raw_group['time_scale'].value

    # Make signal to noise ratio.
    snr_dset = make_dataset(h5, 'snr_spectrum', time_scale.shape)
    spectrum_dset = h5['fft_spectrum_mean']
    if older(snr_dset, [spectrum_dset, raw_group, Timer_object(1434015914)]):
        if verbose > 0:
            print 'Updating the signal to noise ratio.',
            print ' In "update_with_noise_and_response()"',
            print ' with file_name={}'.format(file_name)
        snr_dset[:] = construct_snr_spectrum(h5, plot=plot)
        snr_dset.attrs['time_stamp'] = time.time()

    # Deconvolute the response function
    time_signal_dset = raw_group['time_signal']
    deconv_time_signal_dset = make_dataset(h5, 'filtered_time_signal',
                                           time_signal_dset.shape)
    if older(deconv_time_signal_dset, [raw_group, snr_dset]):
        response, t_response = get_response(plot=plot, verbose=verbose-1)
        if verbose > 0:
            print 'Deconvolving traces.'
            print ' In "update_with_noise_and_response()"',
            print ' with file_name={}'.format(file_name),
            print ' {} events to process.'.format(n_events)
        deconvolver = wiener.Deconcolver(snr_dset.value, response)
        for i_evt in range(n_events):
            deconv_time_signal_dset[i_evt, :] = deconvolver.deconvolve(
                time_signal_dset[i_evt, :])
            update_progress(i_evt, n_events, verbose)
        print ''
        deconv_time_signal_dset.attrs['time_stamp'] = time.time()

    # Calculate the center of mass of the streak peak
    time_com_dset = make_dataset(h5, 'streak_peak_center', (n_events, ))
    photo_line_intensity_dset = make_dataset(h5, 'streak_peak_integral',
                                             (n_events, ))
    if older(time_com_dset, [deconv_time_signal_dset,
                             Timer_object(1443006988)]):
        if verbose > 0:
            print 'Calculating streak peak center in time.',
            print ' In "update_with_noise_and_response()"',
            print ' with file_name={}'.format(file_name)

        streak_sl = slice(np.searchsorted(time_scale, streak_time_roi[0]),
                          np.searchsorted(time_scale, streak_time_roi[1],
                                          side='right'))
        time_scale_streak = time_scale[streak_sl]

        ####
        # Center of mass calculation
#        for i_evt in range(n_events):
#            time_com_dset[i_evt] = get_com(
#                time_scale_streak,
#                deconv_time_signal_dset[i_evt, streak_sl])
#            update_progress(i_evt, n_events, verbose)

        ####
        # Fit of Gaussian
        deconv_time_signal = deconv_time_signal_dset.value
        time_com = np.zeros(time_com_dset.shape)
        photo_line_intensity = np.zeros(photo_line_intensity_dset.shape)
        mean_signal = deconv_time_signal[:, streak_sl].mean(axis=0)

        mod = lmfit.models.GaussianModel()
        params = lmfit.Parameters()
        params.add_many(('amplitude', 1, True, 0),
                        ('center', time_scale_streak[np.argmax(mean_signal)],
                         True, min(time_scale_streak), max(time_scale_streak)),
                        ('sigma', 1e-3, True, 0))
        # fit to mean in order to get start parameters for the shot fits
        out = mod.fit(mean_signal, x=time_scale_streak, params=params)
        for k in params:
            params[k].value = out.params[k].value

        for i_evt in range(n_events):
            out = mod.fit(deconv_time_signal[i_evt, streak_sl],
                          params, x=time_scale_streak)
            time_com[i_evt] = out.params['center'].value
            photo_line_intensity[i_evt] = out.params['amplitude'].value
            update_progress(i_evt, n_events, verbose)

        if plot:
            time_scale_streak = time_scale[streak_sl]
            plt.figure('peak finding time domain')
            plt.clf()
            plt.plot(time_scale_streak, mean_signal)
            plt.plot(time_scale_streak, out.best_fit)

        if verbose > 0:
            print ''
        time_com_dset[:] = time_com
        time_com_dset.attrs['time_stamp'] = time.time()
        photo_line_intensity_dset[:] = photo_line_intensity
        photo_line_intensity_dset.attrs['time_stamp'] = time.time()

    h5.close()


def update_with_time_to_energy_conversion(file_name, plot=False, verbose=0):
    """ Make derived data based on time to energy conversion."""

    update_with_noise_and_response(file_name, plot, verbose)

    h5 = open_hdf5_file(file_name, plot, verbose)
    raw_group = h5['raw']
    n_events = raw_group['event_time_s'].shape[0]

    deconv_time_signal_dset = h5['filtered_time_signal']

    energy_scale_dset = make_dataset(h5, 'energy_scale_eV',
                                     energy_scale_eV.shape)
    energy_trace_dset = make_dataset(h5, 'energy_signal',
                                     (n_events, len(energy_scale_eV)))

    check_tof_to_energy_conversion_matrix(verbose=verbose)
    with h5py.File(tof_to_energy_conversion_file_name, 'r') as tof_to_e_h5:
        if older(energy_scale_dset, [tof_to_e_h5['matrix'],
                                     deconv_time_signal_dset,
                                     Timer_object(1443190000)]):

            if verbose > 0:
                print 'Updating time to energy conversion.',
                print ' In "update_with_time_to_energy_conversion()"',
                print ' with {}'.format(file_name)
            # Get the transformation matrix from file
            M = tof_to_e_h5['matrix'].value
            # Update the energy scale
            energy_scale_dset[:] = tof_to_e_h5['energy_scale_eV'].value
            energy_scale_dset.attrs['time_stamp'] = time.time()
            # Get the photon energy prediction parameters
            params = (tof_to_energy.photon_energy_params() +
                      tof_to_energy.tof_prediction_params())
            for k in params:
                params[k].value = tof_to_e_h5[k].value
            if verbose > 0:
                print 'Computing energy spectra.'
            for i_evt in range(n_events):
                # Energy spectra
                energy_trace_dset[i_evt, :] = M.dot(
                    deconv_time_signal_dset[i_evt, :])
                update_progress(i_evt, n_events, verbose)
            if verbose > 0:
                print ''
            energy_trace_dset.attrs['time_stamp'] = time.time()

    # Calculate energy trace properties
    spectral_properties_group = h5.require_group('spectral_properties')
    spectral_center_dset = make_dataset(spectral_properties_group,
                                        'center_eV', (n_events, ))
    spectral_width_dset = make_dataset(spectral_properties_group,
                                       'width_eV', (n_events, ))
    spectral_threshold_dset = make_dataset(spectral_properties_group,
                                           'threshold', (n_events, ))
    spectral_gaussian_center_dset = make_dataset(spectral_properties_group,
                                                 'gaussian_center',
                                                 (n_events,))

    if older(spectral_center_dset, [energy_trace_dset,
                                    Timer_object(1443421560)]):
        energy_scale = energy_scale_dset[:]
        sl = slice(np.searchsorted(energy_scale, 75),
                   np.searchsorted(energy_scale, 125))
        energy_scale = energy_scale[sl]
        model = lmfit.models.GaussianModel()
        if verbose > 0:
            print 'Calculating spectral center and width:',
            print 'In "update_with_time_to_energy_conversion()"',
            print 'with {}'.format(file_name)
        for i_evt in range(n_events):
            energy_trace = energy_trace_dset[i_evt, sl]
            t_start, t_end, spectral_threshold_dset[i_evt] = \
                get_trace_bounds(energy_scale,
                                 energy_trace,
                                 threshold=8e-5,
                                 min_width=3,
#                                 useRel=True,
#                                 threshold_rel=0.3
                                 )
            center = (t_start + t_end) / 2
            spectral_center_dset[i_evt] = center
            width = t_end - t_start
            spectral_width_dset[i_evt] = width

            # Calculate center of mass
            peak_sl = slice(energy_scale.searchsorted(t_start - width/2),
                            energy_scale.searchsorted(t_end + width/2,
                                                      side='right'))
            peak_trace = energy_trace[peak_sl]
            peak_scale = energy_scale[peak_sl]

#            spectral_com_dset[i_evt] = (np.sum(peak_scale * peak_trace) /
#                                        np.sum(peak_trace))
            if len(peak_trace) > 3:
                out = model.fit(peak_trace, x=peak_scale,
                                center=center, sigma=width/4,
                                amplitude=peak_trace.max() * width / 2)

                spectral_gaussian_center_dset[i_evt] = out.values['center']
            else:
                spectral_gaussian_center_dset[i_evt] = np.nan

            update_progress(i_evt, n_events, verbose)
        spectral_center_dset.attrs['time_stamp'] = time.time()
        spectral_width_dset.attrs['time_stamp'] = time.time()
        spectral_threshold_dset.attrs['time_stamp'] = time.time()
        spectral_gaussian_center_dset.attrs['time_stamp'] = time.time()

    if plot:
        selected_shots = list(np.linspace(0, n_events, 16, endpoint=False))
        plt.figure('peak properties')
        plt.clf()
        _, ax_list = plt.subplots(4, 4, sharex=True, sharey=True,
                                  num='peak properties')
        energy_scale = energy_scale_dset[:]
        sl = slice(np.searchsorted(energy_scale, 75),
                   np.searchsorted(energy_scale, 130))
        energy_scale = energy_scale[sl]
        for i, shot in enumerate(selected_shots):
            energy_trace = energy_trace_dset[shot, :]
            ax = ax_list.flatten()[i]
#            plt.plot(energy_scale - pe_energy_prediction_dset[shot],
            ax.plot(energy_scale, energy_trace[sl])
            c = spectral_center_dset[shot]
            w = spectral_width_dset[shot]
            th = spectral_threshold_dset[shot]
            ax.plot([c-w/2, c+w/2], [th] * 2)

    # Calculate main photoline area
    main_photoline_area = make_dataset(spectral_properties_group,
                                       'main_photoline_area', (n_events, ))

    if older(main_photoline_area, energy_trace_dset):
        if verbose:
            print 'Computing photoline area'
        e_scale = energy_scale_dset.value
        dE = np.mean(np.diff(e_scale))
        e_slice = slice(np.searchsorted(e_scale, 55), None)
        for i_evt in range(n_events):
            raw_A, _ = area_fill.zero_crossing_area(
                energy_trace_dset[i_evt, e_slice])
            main_photoline_area[i_evt] = raw_A * dE
            update_progress(i_evt, n_events, verbose)
        time_stamp_object(main_photoline_area)

    ##########
    # Calculate electron energy prediction
    e_energy_prediction_params_group = make_group(h5,
                                                  'e_energy_prediction_params')

    if older(e_energy_prediction_params_group, [spectral_gaussian_center_dset,
                                                Timer_object(1444931900)]):
        if verbose > 0:
            print 'Fit the electron energy prediction parameters.',
            print 'In "update_with_time_to_energy_conversion()"',
            print 'with {}'.format(file_name)

        selection = np.isfinite(spectral_gaussian_center_dset.value)
#        &
#                     (0.4 < raw_group['phase_cavity_times'][:, 0]) &
#                     (raw_group['phase_cavity_times'][:, 0] < 1.1))
        spectral_gaussian_center = spectral_gaussian_center_dset[selection]
        if len(spectral_gaussian_center) == 0:
            return

        var_dict = {
            'l3_energy': raw_group['energy_L3_MeV'][selection],
            'bc2_energy': h5['energy_BC2_MeV'][selection],
            # 'fee': h5['fee_mean'][selection],
            'e_energy': spectral_gaussian_center
            }
        prediction_params = \
            tof_to_energy.e_energy_prediction_model_start_params(**var_dict)

        try:
            res = lmfit.minimize(tof_to_energy.e_energy_prediction_model,
                                 prediction_params,
                                 kws=var_dict)
            fit_worked = True
        except:
            fit_worked = False

        if verbose > 0 and fit_worked:
            print '\nPrediction params:'
            lmfit.report_fit(res)

        # Create or update the parameters from the fit in the group
        for k, v in prediction_params.iteritems():
            d = e_energy_prediction_params_group.require_dataset(
                k, (), np.float)
            d[()] = v.value if fit_worked else np.nan

        # Remove old parameters that should not be there
        for k in set(e_energy_prediction_params_group.keys()).difference(
                set(prediction_params.keys())):
            del e_energy_prediction_params_group[k]

        e_energy_prediction_params_group.attrs[time_stamp] = time.time()

        if plot:
            deviation = tof_to_energy.e_energy_prediction_model(
                         prediction_params, **var_dict)
            plt.figure('e energy prediction {}'.format(
                h5.filename.split('/')[-1]))
            plt.clf()
            plt.subplot(221)
#            plt.plot(spectral_gaussian_center, deviation, '.')
            plt.scatter(spectral_gaussian_center, deviation,
                        s=4, c=h5['energy_BC2_MeV'][selection],
                        linewidths=(0,), alpha=1)
            plt.xlabel('electron energy (eV)')
            plt.ylabel('prediction residual (eV)')

            x_range = plt.xlim()
            y_range = plt.ylim()
            img, _, _ = np.histogram2d(spectral_gaussian_center, deviation,
                                       bins=2**7, range=[x_range, y_range])
            img = img.T

            plt.subplot(222)
            plt.imshow(img, aspect='auto', interpolation='none',
                       origin='lower', extent=x_range + y_range)

            hist, hist_edges = np.histogram(deviation,
                                            bins=2**5, range=(-3, 3))
            hist_centers = (hist_edges[: -1] + hist_edges[1:])/2
            plt.subplot(223)
            gauss_model = lmfit.models.GaussianModel()
            fit_out = gauss_model.fit(hist, x=hist_centers)
            lmfit.report_fit(fit_out)
            plt.bar(hist_edges[:-1], hist, width=np.diff(hist_edges))
            plt.plot(hist_centers, fit_out.best_fit, 'r', linewidth=2)

            plt.subplot(224)
            plt.plot(spectral_gaussian_center, h5['energy_BC2_MeV'][selection],
                     '.')


def update_with_energy_prediction(file_name, plot=False, verbose=0):

    update_with_time_to_energy_conversion(file_name, plot, verbose)

    h5 = open_hdf5_file(file_name, plot, verbose)
    raw_group = h5['raw']
    n_events = raw_group['event_time_s'].shape[0]

    prediction_map = {'117': 'h5_files/run118_all.h5',
                      '114': 'h5_files/run115_all.h5',
                      '113': 'h5_files/run112_all.h5',
                      '108': 'h5_files/run109_all.h5',
                      '101': 'h5_files/run100_all.h5',
                      '102': 'h5_files/run100_all.h5'}

    pe_energy_prediction_dset = make_dataset(
        h5, 'photoelectron_energy_prediction_eV', (n_events,))
    spectral_properties_group = h5['spectral_properties']
#    spectral_gaussian_center_dset = spectral_properties_group[
#        'gaussian_center']
    fee_dset = h5['fee_mean']
    energy_BC2_dset = h5['energy_BC2_MeV']
    energy_L3_dset = raw_group['energy_L3_MeV']

    for k, v in prediction_map.iteritems():
        if k in file_name:
            update_with_time_to_energy_conversion(v, plot=False,
                                                  verbose=verbose-1)
            ref_h5 = open_hdf5_file(file_name)
            e_energy_prediction_params_group = \
                ref_h5['e_energy_prediction_params']
            break
    else:
        e_energy_prediction_params_group = h5['e_energy_prediction_params']

    if older(pe_energy_prediction_dset, [e_energy_prediction_params_group,
                                         fee_dset,
                                         energy_BC2_dset,
                                         raw_group,
                                         Timer_object(1444981500)]):

        if verbose > 0:
            print 'Updating energy prediction.',
            print ' In "update_with_energy_prediction()" with {}'.format(
                file_name)
        prediction_params = lmfit.Parameters()
        for k in e_energy_prediction_params_group:
            prediction_params.add(k, e_energy_prediction_params_group[k][()])

        var_dict = {
            'l3_energy': energy_L3_dset.value,
            'bc2_energy': energy_BC2_dset.value,
            'fee': fee_dset.value
            }

        try:
            pe_energy_prediction_dset[:] = \
                tof_to_energy.e_energy_prediction_model(prediction_params,
                                                        **var_dict)
        except:
            pe_energy_prediction_dset[:] = np.nan

        pe_energy_prediction_dset.attrs[time_stamp] = time.time()

    ##########
    # Make the christmas three histogram
    n_spectral_center_bins = 2**7
    n_spectral_width_bins = 2**7
    spectral_center_axis_dset = make_dataset(spectral_properties_group,
                                             'center_axis_eV',
                                             (n_spectral_center_bins, ))
    spectral_width_axis_dset = make_dataset(spectral_properties_group,
                                            'width_axis_eV',
                                            (n_spectral_width_bins, ))
    spectral_histogram_dset = make_dataset(spectral_properties_group,
                                           'histogram',
                                           (n_spectral_width_bins,
                                            n_spectral_center_bins))

    spectral_center_dset = spectral_properties_group['center_eV']
    spectral_width_dset = spectral_properties_group['width_eV']
    pct_filter_dset = h5['pct_filter']

    if older(spectral_histogram_dset, [spectral_center_dset,
                                       spectral_width_dset,
                                       pe_energy_prediction_dset,
                                       pct_filter_dset,
                                       Timer_object(2444203160)]):
        if verbose > 0:
            print 'Making the christmas tree plot.',
            print ' In "update_with_energy_prediction()"',
            print ' with {}'.format(file_name)
        spectral_width_axis_dset[:] = np.linspace(0, 35, n_spectral_width_bins)
        spectral_width_axis_dset.attrs['time_stamp'] = time.time()
        spectral_center_axis_dset[:] = np.linspace(-20, 20,
                                                   n_spectral_center_bins)
        spectral_center_axis_dset.attrs['time_stamp'] = time.time()

#        I = (pct_filter_dset.value &
#             (-0.1 < raw_group['phase_cavity_times'][:, 1]) &
##             (raw_group['phase_cavity_times'][:, 1] < 0.05) &
##             (0.75 < raw_group['phase_cavity_times'][:, 0]) &
##             (raw_group['phase_cavity_times'][:, 0] < 0.85) &
#             (0.065 < raw_group['power_meter_V'].value) &
#             (raw_group['power_meter_V'].value < 0.1))

        I = np.ones(pct_filter_dset.shape, dtype=bool)

        hist = aol_plotting.center_histogram_2d(
            spectral_center_dset[I] - pe_energy_prediction_dset[I],
            spectral_width_dset[I],
            spectral_center_axis_dset[:],
            spectral_width_axis_dset[:])
        hist[hist == 0] = np.nan
        spectral_histogram_dset[:] = hist
        spectral_histogram_dset.attrs['time_stamp'] = time.time()

    if plot:
        plt.figure('christmas tree {}'.format(h5.filename.split('/')[-1]))
        plt.clf()
        plt.imshow(spectral_histogram_dset[:], aspect='auto',
                   interpolation='none', origin='lower',
                   extent=(np.min(spectral_center_axis_dset),
                           np.max(spectral_center_axis_dset),
                           np.min(spectral_width_axis_dset),
                           np.max(spectral_width_axis_dset)))

        plt.xlabel('center (eV)')
        plt.ylabel('width (eV)')
        plt.colorbar()
        plt.savefig('figures/christmas_tree_{}.png'.format(
            h5.filename.split('/')[-1].split('.')[0]))

    h5.close()


def load_file(file_name, plot=False, verbose=0):
    """ Load file and make sure it is up to date."""
#    if verbose > 0:
#        print 'Entering "load_file()" with file_name={}'.format(file_name)

    update_with_energy_prediction(file_name, plot, verbose)

    h5 = open_hdf5_file(file_name, plot, verbose)
    raw_group = h5['raw']
    n_events = raw_group['event_time_s'].shape[0]

    if verbose > 0:
        print 'File {} processed.'.format(h5.file)
        print 'It contains', n_events, 'events.'
    if verbose > 1:
        list_hdf5_content(h5)

    return h5


def touch_all_files(verbose=2):
    file_names = ['/'.join([data_dir, f]) for f in os.listdir(data_dir) if
                  f.startswith('run') and f.endswith('_all.h5')]

    for name in file_names:
        load_file(name, verbose=verbose)


if __name__ == '__main__':
    # Parset the command line.
    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--hdf5_file', type=str,
                        default='h5_files/run108_all.h5',
                        help='Path to hdf5 file to process')
    parser.add_argument('--plot', action='store_true',
                        help='Display plots. Default: no plots.')
    parser.add_argument('-v', '--verbose', action='count',
                        help='increase output verbosity')
    args = parser.parse_args()

    # Unpack the parser arguments.
    hdf5_file = args.hdf5_file
    plot = args.plot
    verbose = args.verbose

    # If plotting is requested, ryn pyplot in the interactive mode.
    if plot:
        plt.ion()

    if verbose > 0:
        print 'Get the noise spectrum just to make sure it is up to date.'
    get_nois_spectrum(plot=plot, verbose=verbose)

    # Load the given file.
    if verbose > 0:
        print 'Load the requested file: {}'.format(hdf5_file)
    h5 = load_file(hdf5_file, verbose=verbose, plot=plot)

    # Get the raw group of the file.
    raw_group = h5['raw']

    # Number of events in the file.
    n_events = len(raw_group['event_time_s'])

    # Time trace rellated information.
    raw_time = raw_group['time_scale'].value
    raw_traces_dset = raw_group['time_signal']
    filtered_traces = h5['filtered_time_signal']

    # Pulse energy
    raw_fee_dset = raw_group['FEE_energy_mJ']
    n_fee = raw_fee_dset.shape[1]

    # frequency domain
    freq_axis = h5['fft_freq_axis'].value
    fft_mean = h5['fft_spectrum_mean'].value
    snr = h5['snr_spectrum'].value

    if plot and False:
        if verbose > 0:
            print 'Plotting fee correlations.'
        plt.figure('fee')
        plt.clf()
        ax = None
        for i in range(n_fee):
            for k in range(n_fee):
                ax = plt.subplot(n_fee, n_fee, i + k*n_fee + 1,
                                 sharex=ax, sharey=ax)
                ax.plot(raw_fee_dset[:, i], raw_fee_dset[:, k], '.')
                if i > 0:
                    plt.setp(ax.get_yticklabels(), visible=False)
                if k < n_fee-1:
                    plt.setp(ax.get_xticklabels(), visible=False)
        plt.xlim(xmin=0)
        plt.ylim(ymin=0)

        if verbose > 0:
            print 'Plotting fee histogram.'
        plt.figure('fee histogram')
        plt.clf()
        plt.hist(h5['fee_mean'].value, bins=100)

    if plot:
        if verbose > 0:
            print 'Plot signal maximium histogram.'
        plt.figure('signal hist')
        plt.clf()
        plt.hist(h5['max_signal'], bins=100)

    if plot:
        if verbose > 0:
            print 'Plot spectr'
        plt.figure('fft')
        plt.clf()
        plt.semilogy(freq_axis, fft_mean, label='average spectrum')
        plt.semilogy(freq_axis, snr, label='snr')
        plt.legend(loc='best')

    # Plot some traces
    if plot:
        if verbose > 0:
            print 'Plotting traces'
        trace_fig = plt.figure('traces {}'.format(hdf5_file))
        trace_fig.clf()
        raw_mean_tr = raw_traces_dset.value.mean(0)
        deconv_mean_tr = filtered_traces.value.mean(0)
        rand_event = np.random.randint(n_events)
        response, _ = get_response(plot=False, verbose=verbose)

        plt.plot(raw_time, raw_traces_dset[rand_event, :],
                 label='single trace')
        plt.plot(raw_time, filtered_traces[rand_event, :],
                 label='Deconv single trace')

        plt.plot(raw_time, raw_mean_tr, label='mean trace')
        plt.plot(raw_time, deconv_mean_tr,
                 label='Deconv mean')
        plt.legend(loc='best')

    # Plot the phase cavity times
    pct = raw_group['phase_cavity_times']
    plt.figure('Phase cavity times')
    plt.clf()
#    pc_selection = (np.isfinite(np.sum(pct, axis=1)) &
#                    (pct[:, 0] > -2) & (pct[:, 0] < 2) &
#                    (pct[:, 1] > -2) & (pct[:, 1] < 2))
#                    (pct[:, 0] > -50) & (pct[:, 0] < 50))
    pc_selection = h5['pct_filter'].value

    for i in range(2):
        plt.subplot(1, 3, i+1)
        plt.title('Time {}'.format(i))
        hist, hist_edges = np.histogram(pct[pc_selection, i], bins=100)
        plt.bar(hist_edges[: -1], hist, width=np.diff(hist_edges))

    plt.subplot(133)
    plt.plot(pct[pc_selection, 0], pct[pc_selection, 1], '.')

    # Plot energy traces and photon energy diagnostics
    pe_energy_dset = h5['photoelectron_energy_prediction_eV']
    energy_scale = h5['energy_scale_eV'][:]
    energy_signal_dset = h5['energy_signal']
    selected_shots = np.linspace(0, n_events, 100, endpoint=False, dtype=int)
    plt.figure('Energy spectra')
    plt.clf()
    ax1 = plt.subplot(121)
    ax2 = plt.subplot(122)
    dy = 1e-5
    for i, shot in enumerate(selected_shots):
        ax1.plot(energy_scale, energy_signal_dset[shot, :] + dy * i)
        ax2.plot(energy_scale - pe_energy_dset[shot],
                 energy_signal_dset[shot, :] + dy * i)
    ax2.set_xlim(-20, 25)

# %%
    # Plot the photoline area
    plt.figure('photoline area')
    plt.clf()

    spectral_properties_group = h5['spectral_properties']
    main_photoline_area = spectral_properties_group[
        'main_photoline_area'].value
    fee = h5['fee_mean'].value
    I = np.isfinite(main_photoline_area) & np.isfinite(fee)
    p = np.polyfit(fee[I], main_photoline_area[I], 2)
    fee_ax = np.linspace(min(fee[I]), max(fee[I]), 2**5)

    plt.subplot(121)
    plt.plot(fee, main_photoline_area, '.')
    plt.plot(fee_ax, np.polyval(p, fee_ax), 'r')
    plt.subplot(122)
    plt.hist2d(fee[I], main_photoline_area[I], bins=2**7)
    plt.plot(fee_ax, np.polyval(p, fee_ax), 'r')
