# -*- 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')


# -*- coding: utf-8 -*-
"""
Created on Mon Jun 22 17:15:03 2015

@author: Tim Kunze

Copyright (C) 2015, Tim Kunze. All rights reserved.

This script is a modified version of the RUM Detector:

instead of sweeping over He and Hi in every diagram, we sweep over lenge and intensity of the impulse (as in the actiation plot)

"""

###############################################################################
#
# Imports
#
###############################################################################

import numpy as np
import sys
import scipy as sc
import os               # to enable some C commands (cwd,listdir)

currpath = '/usr/wrk/people9/tiku2449/EI_RUM/001_Unifying_Framework/RUM_Exploration'
os.chdir(currpath)

import sys
sys.path.append("/usr/wrk/people9/tiku2449/EI_RUM/001_Unifying_Framework")

import Models.JuRClass_fin_006 as FCV
import Simulation_And_Analysis.Sim_Simulation_003 as simulate

while len(sys.argv) > 1:
    option = sys.argv[1];                                             del sys.argv[1]
    if   option == '-he':  he = float(sys.argv[1].replace(',','.'));  del sys.argv[1]
    elif option == '-hi':  hi = float(sys.argv[1].replace(',','.'));  del sys.argv[1]
    else:
        print 'Options invalides :',option,'->',sys.argv[0]

#%%
###############################################################################
#
# Main
#
###############################################################################

dt = 1000e-6

JR = FCV.JuR()

JR.integ_stepsize = dt

JR.n=2
JR.coupling = np.array([[0.0,0.0],[0.0,0.0]])   #

JR.distanceMatrix = np.array([[0.0,0.01],[0.0,0.0]]) # important!!
JR.init = np.zeros((8,JR.n))
JR.c_e=0        # only relevant for connected areas
JR.c_i=0        # only relevant for connected areas
JR.c_py=30       # only relevant for connected areas
JR.configure()

#%%
###############################################################################
#
## Activation Diagram RUM with modulation of input to II
#
###############################################################################

t_simulation = 5

N=t_simulation/dt
time = np.arange(0,N*dt,dt)

JR.H_e=he
JR.H_i=hi

p_sim_py  = np.zeros((N,JR.n))
p_sim_e  = np.zeros((N,JR.n))
p_sim_i  = np.zeros((N,JR.n))
length_range = np.arange(500,1501,10)
intensity_range = np.arange(50,251,2)
state_grid = np.zeros((len(intensity_range),len(length_range),3))

i=0
j=0
for ins in intensity_range:
    j=0
    for le in length_range:
        p_sim_e  = np.zeros((N,JR.n))
        p_sim_e[1000:1000+le,:] = ins
        signal,sig_ei,sig_ii,impact,data = simulate.simulate_network_006(JR,p_sim_py,p_sim_e,p_sim_i,t_simulation)
        state_grid[i,j,0] = np.mean(signal[999,0])
        state_grid[i,j,1] = np.mean(signal[4000:,0])
        state_grid[i,j,2] = np.max(signal[900:,0])
        print "len: %.0f | int: %.0f | he: %.2fmV | hi: %2.fmV" %(le, ins, he*1000,hi*1000)
        j+=1

    i+=1

#dataa=length_range,intensity_range,state_grid
np.save('RUM_Dec_meas_full2_le500t1500i10msInt50t250i2_He%.2fmV_Hi%.1fmV.npy' %(he*1000,hi*1000),state_grid)
#np.save('RUM_Dec_sim_le500t1500i10msInt70t250i2_He%.2fmV_Hi%.1fmV.npy' %(he*1000,hi*1000),signal)
#np.save('RUM_Dec_data_le500t1500i10msInt70t250i2_He%.2fmV_Hi%.1fmV.npy' %(he*1000,hi*1000),dataa)

#
#
#def cleargrid(state_grid):
#    [x,y,z]=np.shape(state_grid)
#    for i in range(x):
#        for j in range(y):
#            if state_grid[i,j,1] > 0.004:
#                state_grid[i,j,1] = 0.006
#            elif state_grid[i,j,1] < 0.004:
#                state_grid[i,j,1] = -0.002
#            else:
#                raise ValueError('Error')
#                print "ERROR"
#
#    return state_grid
###
#%% Analysis
#import matplotlib.pyplot as plt
#hirange = np.arange(19,26,1)*1e-3
#herange = np.arange(2.5,4.1,0.25)*1e-3
#
#
#glob_low_val=1e3
#glob_high_val=-1e3
#
#for he in herange:
#    for hi in hirange:
#        a=np.load('RUM_Detector2_Imple500t1500i10msInt70t250i2_He%.2fmV_Hi%.1fmV.npy' %(he*1000,hi*1000))
#
#
#        length_range=a[0]
#        intensity_range=a[1]
#        state_grid=a[2]
#
#        low_lenge=np.min(length_range)
#        high_lenge=np.max(length_range)
#        low_inte=np.min(intensity_range)
#        high_inte=np.max(intensity_range)
#
#        if np.min(state_grid[:,:,1]) < glob_low_val:
#            glob_low_val=np.min(state_grid[:,:,1])
#            print he,hi,glob_low_val
#        if np.max(state_grid[:,:,1]) > glob_high_val:
#            glob_high_val=np.max(state_grid[:,:,1])
#            print he,hi,glob_high_val,1
#
#        plt.figure(2)
#        plt.clf()
#        state_grid=cleargrid(state_grid)
#        plt.imshow(np.flipud(state_grid[:,:,1]), aspect='auto', extent = (low_lenge,high_lenge,low_inte,high_inte),interpolation='none')
#        plt.ylabel('intensity')
#        plt.xlabel('length')
#        plt.title('Detektor Diagram,he:%.0fms, hi:%.0fpps' %(he*1000,hi*1000))
#        cb=plt.colorbar()
#        plt.savefig('RUM_Detektor2_Imple500t1500i10msInt70t250i2_He%.2fmV_Hi%.1fmV.pdf'  %(he*1000,hi*1000), format='pdf', dpi=1000)
#        plt.close()
#        #
#
#        # baselevel plot hier zwecklos, da baselevel bei allen stimuli gleich
#        plt.figure(2)
#        plt.clf()
#        #state_grid=cleargrid(state_grid)
#        plt.clf()
#        plt.imshow(np.flipud(state_grid[:,:,0]), aspect='auto', extent = (low_lenge,high_lenge,low_inte,high_inte),interpolation='none')
#        plt.ylabel('intensity')
#        plt.xlabel('length')
#        plt.title('Detektor Diagram,Baselevels,he:%.0fmV, hi:%.0fmV' %(he*1000,hi*1000))
#        plt.colorbar()
#        #plt.savefig('RUM_Detektor_Baselevel_Imple%.0fmsInt%.0f_He2.5t7.0i0k05_Hi10t25i0k1_1.pdf' %(lenge,inte), format='pdf', dpi=1000)
#
#        plt.close('all')


import gym
import os
import sys
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from gym import wrappers
from datetime import datetime
from sklearn.pipeline import FeatureUnion
from sklearn.preprocessing import StandardScaler
from sklearn.kernel_approximation import RBFSampler
from sklearn.linear_model import SGDRegressor

class FeatureTransformer:
	def __init__(self, env, n_components=500):
		observation_examples = np.array([env.observation_space.sample() for x in range(1000)])
		scaler = StandardScaler()
		scaler.fit(observation_examples)

		featurizer = FeatureUnion([
			("rbf1", RBFSampler(gamma=5.0, n_components=n_components)),
			("rbf2", RBFSampler(gamma=4.0, n_components=n_components)),
			("rbf3", RBFSampler(gamma=3.0, n_components=n_components)),
			("rbf4", RBFSampler(gamma=2.0, n_components=n_components)),
			("rbf5", RBFSampler(gamma=1.0, n_components=n_components)),
			("rbf6", RBFSampler(gamma=0.5, n_components=n_components)),
		])
		example_features = featurizer.fit_transform(scaler.transform(observation_examples))

		self.dimensions = example_features.shape[1]
		self.scaler = scaler
		self.featurizer = featurizer

	def transform(self, observation):
		scaled = self.scaler.transform(observation)
		return self.featurizer.transform(scaled)

class Model:
	def __init__(self, env, feature_transformer, learning_rate):
		self.env = env
		self.models = []
		self.feature_transformer = feature_transformer
		for i in range(env.action_space.n):
			model = SGDRegressor(learning_rate)
			model.partial_fit(feature_transformer.transform([env.reset()]), [0])
			self.models.append(model)

	def predict(self, s):
		X = self.feature_transformer.transform([s])
		assert(len(X.shape) == 2)
		return np.array([m.predict(X)[0] for m in self.models])

	def update(self, s, a, G):
		X = self.feature_transformer.transform([s])
		assert(len(X.shape) == 2)
		self.models[a].partial_fit(X, [G])

	def sample_action(self, s, eps):
		if np.random.random() < eps:
			return self.env.action_space.sample()
		else:
			return np.argmax(self.predict(s))

def play_one(model, eps, gamma):
	observation = env.reset()
	done = False
	totalreward = 0
	iters = 0
	while not done and iters < 1000:
		action = model.sample_action(observation, eps)
		prev_observation = observation
		observation, reward, done, info = env.step(action)

		# Update the model
		G = reward + gamma*np.max(model.predict(observation)[0])
		model.update(prev_observation, action, G)

		totalreward += reward
		iters += 1

	return totalreward

def plot_cost_to_go(env, estimator, num_tiles=20):
	x = np.linspace(env.observation_space.low[0], env.observation_space.high[0], num=num_tiles)
	y = np.linspace(env.observation_space.low[1], env.observation_space.high[1], num=num_tiles)
	X, Y = np.meshgrid(x, y)

	Z = np.apply_along_axis(lambda _: -np.max(estimator.predict(_)), 2, np.dstack([X, Y]))

	fig = plt.figure(figsize=(10, 5))
	ax = fig.add_subplot(111, projection='3d')
	surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=matplotlib.cm.coolwarm, vmin=-1.0, vmax=1.0)
	ax.set_xlabel('Position')
	ax.set_ylabel('Velocity')
	ax.set_zlabel('Cost-To-Go == -V(s)')
	ax.set_title("Cost-To-Go Function")
	fig.colorbar(surf)
	plt.show()

def plot_running_avg(totalrewards):
	N = len(totalrewards)
	running_avg = np.empty(N)
	for t in range(N):
		running_avg[t] = totalrewards[max(0, t-100):(t+1)].mean()
	plt.plot(running_avg)
	plt.title("Running Average")
	plt.show()

def main():
	env = gym.make('MountainCar-v0')
	ft = FeatureTransformer(env)
	model = Model(env, ft, "constant")

	gamma = 0.99

	if 'monitor' in sys.argv:
		filename = os.path.basename(__file__).split('.')[0]
		monitor_dir = './' + filename + '_' + str(datetime.now())
		env = wrappers.Monitor(env, monitor_dir)

	N = 300
	totalrewards = np.empty(N)
	for n in range(N):
		eps = 0.1*(0.97**n)
		totalreward = play_one(model, eps, gamma)
		totalrewards[n] = totalreward
		print("episode:", n, "total reward:", totalreward)
	print("avg reward for last 100 episodes:", totalrewards[-100:].mean())
	print("total steps:", -totalrewards.sum())

	plt.plot(totalrewards)
	plt.title("Rewards")
	plt.show()

	plot_running_avg(totalrewards)

	plot_cost_to_go(env, model)

if __name__ == '__main__':
	main()


#import yt.mods as yt
import yt
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import glob

__all__ = ['future_snr', 'snr']

_core_collapse_labels = ["SNII", "II", "2", "SN_II", "TypeII", "Type 2",
                         "Type II", "type II", "typeII", 'core collapse']

_snia_labels = ["SN1a", "SNIa", "Type1a", "TypeIa", "Type Ia", "Type 1a",
                 "type 1a", "type Ia", "type ia", "type1a", "typeIa"]

_agb_labels = ['AGB', 'agb']

def future_snr(ds, data, times = None, sn_type = 'II'):
    """
    Looks forward from current time to compute future (projected) SN
    rate
    """

    current_time  = ds.current_time.convert_to_units('Myr').value

    if times is None:
        bin_spacing = 2.0* yt.units.Myr
        times = np.arange(np.min(creation_time) - bin_spacing*2.0, currentTime, bin_spacing)*yt.units.Myr
    elif np.size(times) == 1:
        bin_spacing = times
        if not hasattr(bin_spacing, 'value'):
            bin_spacing = bin_spacing * yt.units.Myr

        times = np.linspace(current_time, current_time + 2000.0, bin_spacing)
        times = times * yt.units.Myr

    birth_mass    = data['birth_mass'].value
    mass          = data['particle_mass'].convert_to_units('Msun').value
    creation_time = data['creation_time'].convert_to_units('Myr').value
    lifetimes     = data['dynamical_time'].convert_to_units('Myr').value
    pt            = data['particle_type']

    if  any( [sn_type in x for x in _core_collapse_labels]):

        collapse_threshold = ds.parameters['IndividualStarDirectCollapseThreshold']
        agb_threshold      = ds.parameters['IndividualStarSNIIMassCutoff']

        pcut = (pt == 11) * (birth_mass <= collapse_threshold) *\
                            (birth_mass  > agb_threshold)

    elif any( [sn_type in x for x in _snia_labels]):

        pcut = (pt == 12) * (mass > 0.0)

    elif any( [sn_type in x for x in _agb_labels]):

        pcut = (pt == 11) * (birth_mass < agb_threshold)

    explosion_times = creation_time[pcut] + lifetimes[pcut]
    explosion_times = explosion_times * yt.units.Myr

    times = times.convert_to_units('yr')
    snr   = np.zeros(np.size(times.value) - 1)

    # compute SNR
    for i in np.arange(np.size(times) - 1):
        dt = times[i+1] - times[i]
        dN = np.size( explosion_times[explosion_times <= times[i+1]]) -\
             np.size( explosion_times[explosion_times <= times[i]])

        snr[i] = dN / dt

    return times, snr

def snr(ds, data, times = None, sn_type = 'II'):
    """
    Computes the supernova rate of the desired time for a given dataset
    as a function of time. The way the particle types and particle lifetimes
    are handled, this can be done for the entire galaxy history using a single
    snapshot, rather than having to sort through each dump.

    One can provide sample times using "times" argument, or leave it alone for
    a 10 Myr sample spacing from t = 0 to t = current_time. If a single value
    is provided, this is taken to be the sample spacing (dt), sampled over
    t = 0 to t = current_time. Units are assumed to be Myr if not provided.

    Accounts for direct collapse model in computing SNII rates using
    parameter file.
    """

    current_time  = ds.current_time.convert_to_units('Myr').value

    if times is None:
        bin_spacing = 10.0 * yt.units.Myr
        times = np.linspace(np.min(creation_time), current_time, bin_spacing)*yt.units.Myr
    elif np.size(times) == 1:
        bin_spacing = times
        if not hasattr(bin_spacing, 'value'):
            bin_spacing = bin_spacing * yt.units.Myr

        times = np.linspace(np.min(creation_time), current_time, bin_spacing)
        times = times *yt.units.Myr

    # load particle properties
    birth_mass    = data['birth_mass'].value
    mass          = data['particle_mass'].convert_to_units("Msun").value
    creation_time = data['creation_time'].convert_to_units('Myr').value
    metallicity   = data['metallicity_fraction'].value
#    lifetimes     = data['dynamical_time'].convert_to_units('Myr').value
    lifetimes     = data[('io','particle_model_lifetime')].convert_to_units('Myr').value
    pt            = data['particle_type'].value

    # check to see if there are any SN candidates in the first place
#    if not any([ == x for x in np.unique(pt)]):
#        print "no supernova of type " + sn_type + " found"
#        return times, np.zeros(np.size(times.value) - 1)

    # looking for core collapse supernova rate
    if  any( [sn_type in x for x in _core_collapse_labels]):

        pcut = (pt == 13)

        # ignore stars that did not actually go supernova
        collapse_threshold = ds.parameters['IndividualStarDirectCollapseThreshold']
        agb_threshold      = ds.parameters['IndividualStarSNIIMassCutoff']
        if not any([(x <= collapse_threshold)*(x > agb_threshold) for x in birth_mass[pcut]]):
            print("no core collapse supernova present, only direct collapse")
            return times, np.zeros(np.size(times.value) - 1)

        # slice!
        pcut *= (birth_mass <= collapse_threshold)*(birth_mass > agb_threshold)

    elif any( [sn_type in x for x in _snia_labels]):

        pcut = (pt == 12)

        if np.size(mass[pcut]) < 1:
            return times, np.zeros(np.size(times))

        # SNIa are the ones that are just masless tracers, rest are WD
        if not any(mass[pcut] == 0.0):
            print("no Type Ia supernova, only white dwarfs")
            print("N_WD = %i -- Lowest mass = %.3f Msun"%(np.size(mass[pcut]), np.min(mass[pcut])))
            print("Current time = %.2E Myr - Next to explode at t = %.2E Myr"%(current_time, np.min(lifetimes[pcut] + creation_time[pcut])))
            return times, np.zeros(np.size(times.value) - 1)

        # slice!
        pcut *= (mass == 0.0)

    elif any( [sn_type in x for x in _agb_labels]):
        agb_threshold	   = ds.parameters['IndividualStarSNIIMassCutoff']

        pcut = (pt > 11)  # all dead stars
        pcut = pcut * (birth_mass <= agb_threshold)

 #       pcut  = (pt == 12)
 #       pcut *= (mass > 0.0)

  #      pcut = pcut + ( (pt == 13) * (birth_mass <= agb_threshold))

    else:
        print("sn_type :" + sn_type + " not a valid option - check spelling")
        return -1

    #
    # now get the explosion times for all supernova
    # when stars go SN, lifetime is set to be lifetime*huge_number
    # therefore, explosion time can be backed out as:
    #
    explosion_times = creation_time[pcut] + lifetimes[pcut]/ds.parameters['huge_number']
    explosion_times = explosion_times * yt.units.Myr

    times = times.convert_to_units('yr')
    snr   = np.zeros(np.size(times.value) - 1)

    # compute SNR
    for i in np.arange(np.size(times) - 1):
        dt = times[i+1] - times[i]
        dN = np.size( explosion_times[explosion_times <= times[i+1]]) -\
             np.size( explosion_times[explosion_times <= times[i]])

        snr[i] = dN / dt

    return times, snr

if __name__ == '__main__':
    # example usage - uses most recent data file

    log = False

    ds_list = np.sort( glob.glob('./DD????/DD????'))

    ds   = yt.load(ds_list[-1])
    data = ds.all_data()

    dt = 25.0
    times = np.arange(0.0, ds.current_time.convert_to_units('Myr').value + dt, dt)
    times = times*yt.units.Myr

    times, snrII = snr(ds, data, times = times, sn_type = 'TypeII')
    times, snrIa = snr(ds, data, times = times, sn_type = "TypeIa")

    center = 0.5 * (times[1:] + times[:-1])

    fig, ax = plt.subplots(figsize=(8,8))

    snialabel = 'Type Ia x 10'
    sniilabel = 'Core Collapse'

    ftimes = np.arange(ds.current_time.convert_to_units('Myr').value,
                             ds.current_time.convert_to_units('Myr').value + 800.0 + 10, 10)
    ftimes = ftimes * yt.units.Myr
    ftimes, fsnrII = future_snr(ds, data, times = ftimes, sn_type = 'TypeII')
    ftimes, fsnrIa = future_snr(ds, data, times = ftimes, sn_type = 'TypeIa')

    if log:
        ax.plot(center/1.0E6, snrII*1.0E6, color = 'black', lw = 3, ls = '-', label = sniilabel)
        ax.plot(center/1.0E6, snrIa*1.0E6*10, color = 'black', lw = 3, ls = '--', label = snialabel)
        x.semilogy()
    else:
        ax.step(times[:-1]/1.0E6, snrII*1.0E6, color ='black', lw = 3, ls = '-', label = sniilabel)
        ax.step(times[:-1]/1.0E6, snrIa*1.0E6 * 10, color ='orange', lw = 3, ls = '-', label = snialabel)
        ax.step(ftimes[:-1]/1.0E6, fsnrII*1.0E6, color = 'black', lw = 3, ls = ':')
        ax.step(ftimes[:-1]/1.0E6, fsnrIa*1.0E6 * 10, color = 'orange', lw = 3, ls = ':')

    ax.set_xlabel('Time (Myr)')

    ax.set_ylabel(r'SNR (Myr$^{-1}$)')
    ax.set_ylim( np.min( [np.min(snrIa), np.min(snrII)])*1.0E6,
                 np.max( [np.max(snrIa), np.max(snrII)])*1.25*1.0E6)
    ax.plot( [ds.current_time.convert_to_units('Myr').value]*2, ax.get_ylim(), ls = '--', lw = 3, color = 'black')

    ax.legend(loc ='best')
    plt.tight_layout()
    ax.minorticks_on()
    plt.savefig('snr.png')


from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import os
import shutil
import tempfile
from contextlib import contextmanager

import matplotlib as mpl
from matplotlib import style
from matplotlib.style.core import USER_LIBRARY_PATHS, STYLE_EXTENSION

import six

PARAM = 'image.cmap'
VALUE = 'pink'
DUMMY_SETTINGS = {PARAM: VALUE}

@contextmanager
def temp_style(style_name, settings=None):
    """Context manager to create a style sheet in a temporary directory."""
    settings = DUMMY_SETTINGS
    temp_file = '%s.%s' % (style_name, STYLE_EXTENSION)

    # Write style settings to file in the temp directory.
    tempdir = tempfile.mkdtemp()
    with open(os.path.join(tempdir, temp_file), 'w') as f:
        for k, v in six.iteritems(settings):
            f.write('%s: %s' % (k, v))

    # Add temp directory to style path and reload so we can access this style.
    USER_LIBRARY_PATHS.append(tempdir)
    style.reload_library()

    try:
        yield
    finally:
        shutil.rmtree(tempdir)
        style.reload_library()

def test_available():
    with temp_style('_test_', DUMMY_SETTINGS):
        assert '_test_' in style.available

def test_use():
    mpl.rcParams[PARAM] = 'gray'
    with temp_style('test', DUMMY_SETTINGS):
        with style.context('test'):
            assert mpl.rcParams[PARAM] == VALUE

def test_use_url():
    with temp_style('test', DUMMY_SETTINGS):
        with style.context('https://gist.github.com/adrn/6590261/raw'):
            assert mpl.rcParams['axes.facecolor'] == "#adeade"

def test_context():
    mpl.rcParams[PARAM] = 'gray'
    with temp_style('test', DUMMY_SETTINGS):
        with style.context('test'):
            assert mpl.rcParams[PARAM] == VALUE
    # Check that this value is reset after the exiting the context.
    assert mpl.rcParams[PARAM] == 'gray'

if __name__ == '__main__':
    from numpy import testing
    testing.run_module_suite()


# (C) British Crown Copyright 2014 - 2017, Met Office
#
# This file is part of cartopy.
#
# cartopy is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the
# Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# cartopy 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 Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with cartopy.  If not, see <https://www.gnu.org/licenses/>.

from __future__ import (absolute_import, division, print_function)

try:
    from unittest.mock import Mock
except ImportError:
    from mock import Mock
from matplotlib.axes import Axes
import pytest

import cartopy.crs as ccrs
from cartopy.mpl.geoaxes import GeoAxes
from cartopy.mpl.ticker import LatitudeFormatter, LongitudeFormatter

def test_LatitudeFormatter_bad_axes():
    formatter = LatitudeFormatter()
    formatter.axis = Mock(axes=Mock(Axes, projection=ccrs.PlateCarree()))
    message = 'This formatter can only be used with cartopy axes.'
    with pytest.raises(TypeError, message=message):
        formatter(0)

def test_LatitudeFormatter_bad_projection():
    formatter = LatitudeFormatter()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=ccrs.Orthographic()))
    message = 'This formatter cannot be used with non-rectangular projections.'
    with pytest.raises(TypeError, message=message):
        formatter(0)

def test_LongitudeFormatter_bad_axes():
    formatter = LongitudeFormatter()
    formatter.axis = Mock(axes=Mock(Axes, projection=ccrs.PlateCarree()))
    message = 'This formatter can only be used with cartopy axes.'
    with pytest.raises(TypeError, message=message):
        formatter(0)

def test_LongitudeFormatter_bad_projection():
    formatter = LongitudeFormatter()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=ccrs.Orthographic()))
    message = 'This formatter cannot be used with non-rectangular projections.'
    with pytest.raises(TypeError, message=message):
        formatter(0)

def test_LatitudeFormatter():
    formatter = LatitudeFormatter()
    p = ccrs.PlateCarree()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-90, -60, -30, 0, 30, 60, 90]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'90\u00B0S', u'60\u00B0S', u'30\u00B0S', u'0\u00B0',
                u'30\u00B0N', u'60\u00B0N', u'90\u00B0N']
    assert result == expected

def test_LatitudeFormatter_degree_symbol():
    formatter = LatitudeFormatter(degree_symbol='')
    p = ccrs.PlateCarree()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-90, -60, -30, 0, 30, 60, 90]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'90S', u'60S', u'30S', u'0',
                u'30N', u'60N', u'90N']
    assert result == expected

def test_LatitudeFormatter_number_format():
    formatter = LatitudeFormatter(number_format='.2f')
    p = ccrs.PlateCarree()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-90, -60, -30, 0, 30, 60, 90]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'90.00\u00B0S', u'60.00\u00B0S', u'30.00\u00B0S',
                u'0.00\u00B0', u'30.00\u00B0N', u'60.00\u00B0N',
                u'90.00\u00B0N']
    assert result == expected

def test_LatitudeFormatter_mercator():
    formatter = LatitudeFormatter()
    p = ccrs.Mercator()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-15496570.739707904, -8362698.548496634,
                  -3482189.085407435, 0.0, 3482189.085407435,
                  8362698.548496634, 15496570.739707898]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'80\u00B0S', u'60\u00B0S', u'30\u00B0S', u'0\u00B0',
                u'30\u00B0N', u'60\u00B0N', u'80\u00B0N']
    assert result == expected

def test_LatitudeFormatter_small_numbers():
    formatter = LatitudeFormatter(number_format='.7f')
    p = ccrs.PlateCarree()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [40.1275150, 40.1275152, 40.1275154]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'40.1275150\u00B0N', u'40.1275152\u00B0N',
                u'40.1275154\u00B0N']
    assert result == expected

def test_LongitudeFormatter_central_longitude_0():
    formatter = LongitudeFormatter(dateline_direction_label=True)
    p = ccrs.PlateCarree()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-180, -120, -60, 0, 60, 120, 180]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'180\u00B0W', u'120\u00B0W', u'60\u00B0W', u'0\u00B0',
                u'60\u00B0E', u'120\u00B0E', u'180\u00B0E']
    assert result == expected

def test_LongitudeFormatter_central_longitude_180():
    formatter = LongitudeFormatter(zero_direction_label=True)
    p = ccrs.PlateCarree(central_longitude=180)
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-180, -120, -60, 0, 60, 120, 180]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'0\u00B0E', u'60\u00B0E', u'120\u00B0E', u'180\u00B0',
                u'120\u00B0W', u'60\u00B0W', u'0\u00B0W']
    assert result == expected

def test_LongitudeFormatter_central_longitude_120():
    formatter = LongitudeFormatter()
    p = ccrs.PlateCarree(central_longitude=120)
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-180, -120, -60, 0, 60, 120, 180]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'60\u00B0W', u'0\u00B0', u'60\u00B0E', u'120\u00B0E',
                u'180\u00B0', u'120\u00B0W', u'60\u00B0W']
    assert result == expected

def test_LongitudeFormatter_degree_symbol():
    formatter = LongitudeFormatter(degree_symbol='',
                                   dateline_direction_label=True)
    p = ccrs.PlateCarree()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-180, -120, -60, 0, 60, 120, 180]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'180W', u'120W', u'60W', u'0', u'60E', u'120E', u'180E']
    assert result == expected

def test_LongitudeFormatter_number_format():
    formatter = LongitudeFormatter(number_format='.2f',
                                   dateline_direction_label=True)
    p = ccrs.PlateCarree()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-180, -120, -60, 0, 60, 120, 180]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'180.00\u00B0W', u'120.00\u00B0W', u'60.00\u00B0W',
                u'0.00\u00B0', u'60.00\u00B0E', u'120.00\u00B0E',
                u'180.00\u00B0E']
    assert result == expected

def test_LongitudeFormatter_mercator():
    formatter = LongitudeFormatter(dateline_direction_label=True)
    p = ccrs.Mercator()
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-20037508.342783064, -13358338.895188706,
                  -6679169.447594353, 0.0, 6679169.447594353,
                  13358338.895188706, 20037508.342783064]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'180\u00B0W', u'120\u00B0W', u'60\u00B0W', u'0\u00B0',
                u'60\u00B0E', u'120\u00B0E', u'180\u00B0E']
    assert result == expected

def test_LongitudeFormatter_small_numbers_0():
    formatter = LongitudeFormatter(number_format='.7f')
    p = ccrs.PlateCarree(central_longitude=0)
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-17.1142343, -17.1142340, -17.1142337]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'17.1142343\u00B0W', u'17.1142340\u00B0W',
                u'17.1142337\u00B0W']
    assert result == expected

def test_LongitudeFormatter_small_numbers_180():
    formatter = LongitudeFormatter(zero_direction_label=True,
                                   number_format='.7f')
    p = ccrs.PlateCarree(central_longitude=180)
    formatter.axis = Mock(axes=Mock(GeoAxes, projection=p))
    test_ticks = [-17.1142343, -17.1142340, -17.1142337]
    result = [formatter(tick) for tick in test_ticks]
    expected = [u'162.8857657\u00B0E', u'162.8857660\u00B0E',
                u'162.8857663\u00B0E']
    assert result == expected


# encoding: utf-8
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Categorical tests."""

#  limitations under the License.
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

from tensorflow.contrib.learn.python.learn.learn_io import HAS_PANDAS
from tensorflow.contrib.learn.python.learn.preprocessing import categorical

class CategoricalTest(tf.test.TestCase):
  """Categorical tests."""

  def testSingleCategoricalProcessor(self):
    cat_processor = categorical.CategoricalProcessor(min_frequency=1)
    x = cat_processor.fit_transform([["0"], [1], [float("nan")], ["C"], ["C"],
                                     [1], ["0"], [np.nan], [3]])
    self.assertAllEqual(list(x), [[2], [1], [0], [3], [3], [1], [2], [0], [0]])

  def testSingleCategoricalProcessorPandasSingleDF(self):
    if HAS_PANDAS:
      import pandas as pd  # pylint: disable=g-import-not-at-top
      cat_processor = categorical.CategoricalProcessor()
      data = pd.DataFrame({"Gender": ["Male", "Female", "Male"]})
      x = list(cat_processor.fit_transform(data))
      self.assertAllEqual(list(x), [[1], [2], [1]])

  def testMultiCategoricalProcessor(self):
    cat_processor = categorical.CategoricalProcessor(min_frequency=0,
                                                     share=False)
    x = cat_processor.fit_transform([["0", "Male"], [1, "Female"], ["3", "Male"]
                                    ])
    self.assertAllEqual(list(x), [[1, 1], [2, 2], [3, 1]])

if __name__ == "__main__":
  tf.test.main()


"""This file contains code used in "Think Stats",
by Allen B. Downey, available from greenteapress.com

Copyright 2014 Allen B. Downey
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html
"""

from __future__ import print_function

import numpy as np
import pandas

import hinc
import thinkplot
import thinkstats2

"""This file contains a solution to an exercise in Think Stats:

The distributions of wealth and income are sometimes modeled using
lognormal and Pareto distributions.  To see which is better, let's
look at some data.

The Current Population Survey (CPS) is joint effort of the Bureau
of Labor Statistics and the Census Bureau to study income and related
variables.  Data collected in 2013 is available from
http://www.census.gov/hhes/www/cpstables/032013/hhinc/toc.htm.
I downloaded hinc06.xls, which is an Excel spreadsheet with
information about household income, and converted it to hinc06.csv,
a CSV file you will find in the repository for this book.  You
will also find hinc.py, which reads the CSV file.

Extract the distribution of incomes from this dataset.  Are any of the
analytic distributions in this chapter a good model of the data?  A
solution to this exercise is in hinc_soln.py.

My solution generates three figures:

1) The CDF of income on a linear scale.

2) The CCDF on a log-log scale along with a Pareto model intended
to match the tail behavior.

3) The CDF on a log-x scale along with a lognormal model chose to
match the median and inter-quartile range.

My conclusions based on these figures are:

1) The Pareto model is probably a reasonable choice for the top
   10-20% of incomes.

2) The lognormal model captures the shape of the distribution better,
   but the data deviate substantially from the model.  With different
   choices for sigma, you could match the upper or lower tail, but not
   both at the same time.

In summary I would say that neither model captures the whole distribution,
so you might have to

1) look for another analytic model,
2) choose one that captures the part of the distribution that is most
   relevent, or
3) avoid using an analytic model altogether.

"""

class SmoothCdf(thinkstats2.Cdf):
    """Represents a CDF based on calculated quantiles.
    """
    def Render(self):
        """Because this CDF was not computed from a sample, it
        should not be rendered as a step function.
        """
        return self.xs, self.ps

    def Prob(self, x):
        """Compute CDF(x), interpolating between known values.
        """
        return np.interp(x, self.xs, self.ps)

    def Value(self, p):
        """Compute inverse CDF(x), interpolating between probabilities.
        """
        return np.interp(p, self.ps, self.xs)

def MakeFigures(df):
    """Plots the CDF of income in several forms.
    """
    xs, ps = df.income.values, df.ps.values
    cdf = SmoothCdf(xs, ps, label='data')
    cdf_log = SmoothCdf(np.log10(xs), ps, label='data')

    # linear plot
    thinkplot.Cdf(cdf)
    thinkplot.Save(root='hinc_linear',
                   xlabel='household income',
                   ylabel='CDF')

    # pareto plot
    # for the model I chose parameters by hand to fit the tail
    xs, ys = thinkstats2.RenderParetoCdf(xmin=55000, alpha=2.5,
                                         low=0, high=250000)
    thinkplot.Plot(xs, 1-ys, label='model', color='0.8')

    thinkplot.Cdf(cdf, complement=True)
    thinkplot.Save(root='hinc_pareto',
                   xlabel='log10 household income',
                   ylabel='CCDF',
                   xscale='log',
                   yscale='log')

    # lognormal plot
    # for the model I estimate mu and sigma using
    # percentile-based statistics
    median = cdf_log.Percentile(50)
    iqr = cdf_log.Percentile(75) - cdf_log.Percentile(25)
    std = iqr / 1.349

    # choose std to match the upper tail
    std = 0.35
    print(median, std)

    xs, ps = thinkstats2.RenderNormalCdf(median, std, low=3.5, high=5.5)
    thinkplot.Plot(xs, ps, label='model', color='0.8')

    thinkplot.Cdf(cdf_log)
    thinkplot.Save(root='hinc_normal',
                   xlabel='log10 household income',
                   ylabel='CDF')

def main():
    df = hinc.ReadData()
    MakeFigures(df)

if __name__ == "__main__":
    main()


#
# DX Analytics
# Helper Function for Plotting
# dx_plot.py
#
# DX Analytics is a financial analytics library, mainly for
# derviatives modeling and pricing by Monte Carlo simulation
#
# (c) Dr. Yves J. Hilpisch
# The Python Quants GmbH
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or any later version.
#
# This program 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see http://www.gnu.org/licenses/.
#

import matplotlib as mpl; mpl.use('agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from pylab import cm

def plot_option_stats(s_list, pv, de, ve):
    ''' Plot option prices, deltas and vegas for a set of
    different initial values of the underlying.

    Parameters
    ==========
    s_list : array or list
        set of intial values of the underlying
    pv : array or list
        present values
    de : array or list
        results for deltas
    ve : array or list
        results for vega
    '''
    plt.figure(figsize=(9, 7))
    sub1 = plt.subplot(311)
    plt.plot(s_list, pv, 'ro', label='Present Value')
    plt.plot(s_list, pv, 'b')
    plt.grid(True)
    plt.legend(loc=0)
    plt.setp(sub1.get_xticklabels(), visible=False)
    sub2 = plt.subplot(312)
    plt.plot(s_list, de, 'go', label='Delta')
    plt.plot(s_list, de, 'b')
    plt.grid(True)
    plt.legend(loc=0)
    plt.setp(sub2.get_xticklabels(), visible=False)
    sub3 = plt.subplot(313)
    plt.plot(s_list, ve, 'yo', label='Vega')
    plt.plot(s_list, ve, 'b')
    plt.xlabel('Strike')
    plt.grid(True)
    plt.legend(loc=0)

def plot_option_stats_full(s_list, pv, de, ve, th, rh, ga):
    ''' Plot option prices, deltas and vegas for a set of
    different initial values of the underlying.

    Parameters
    ==========
    s_list : array or list
        set of intial values of the underlying
    pv : array or list
        present values
    de : array or list
        results for deltas
    ve : array or list
        results for vega
    th : array or list
        results for theta
    rh : array or list
        results for rho
    ga : array or list
        results for gamma
    '''
    plt.figure(figsize=(10, 14))
    sub1 = plt.subplot(611)
    plt.plot(s_list, pv, 'ro', label='Present Value')
    plt.plot(s_list, pv, 'b')
    plt.grid(True)
    plt.legend(loc=0)
    plt.setp(sub1.get_xticklabels(), visible=False)
    sub2 = plt.subplot(612)
    plt.plot(s_list, de, 'go', label='Delta')
    plt.plot(s_list, de, 'b')
    plt.grid(True)
    plt.legend(loc=0)
    plt.setp(sub2.get_xticklabels(), visible=False)
    sub3 = plt.subplot(613)
    plt.plot(s_list, ve, 'yo', label='Gamma')
    plt.plot(s_list, ve, 'b')
    plt.grid(True)
    plt.legend(loc=0)
    sub4 = plt.subplot(614)
    plt.plot(s_list, th, 'mo', label='Vega')
    plt.plot(s_list, th, 'b')
    plt.grid(True)
    plt.legend(loc=0)
    sub5 = plt.subplot(615)
    plt.plot(s_list, rh, 'co', label='Theta')
    plt.plot(s_list, rh, 'b')
    plt.grid(True)
    plt.legend(loc=0)
    sub6 = plt.subplot(616)
    plt.plot(s_list, ga, 'ko', label='Rho')
    plt.plot(s_list, ga, 'b')
    plt.xlabel('Strike')
    plt.grid(True)
    plt.legend(loc=0)

def plot_greeks_3d(inputs, labels):
    ''' Plot Greeks in 3d.

    Parameters
    ==========
    inputs : list of arrays
        x, y, z arrays
    labels : list of strings
        labels for x, y, z
    '''
    x, y, z = inputs
    xl, yl, zl = labels
    fig = plt.figure(figsize=(10, 7))
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(x, y, z, rstride=1, cstride=1,
                           cmap=cm.coolwarm, linewidth=0.5, antialiased=True)
    ax.set_xlabel(xl)
    ax.set_ylabel(yl)
    ax.set_zlabel(zl)
    fig.colorbar(surf, shrink=0.5, aspect=5)

def plot_calibration_results(cali, relative=False):
    ''' Plot calibration results.

    Parameters
    ==========
    cali : instance of calibration class
        instance has to have opt_parameters
    relative : boolean
        if True, then relative error reporting
        if False, absolute error reporting
    '''
    cali.update_model_values()
    mats = set(cali.option_data[:, 0])
    mats = np.sort(list(mats))
    fig, axarr = plt.subplots(len(mats), 2, sharex=True)
    fig.set_size_inches(8, 12)
    fig.subplots_adjust(wspace=0.2, hspace=0.2)
    z = 0
    for T in mats:
        strikes = strikes = cali.option_data[cali.option_data[:, 0] == T][:, 1]
        market = cali.option_data[cali.option_data[:, 0] == T][:, 2]
        model = cali.model_values[cali.model_values[:, 0] == T][:, 2]
        axarr[z, 0].set_ylabel('%s' % str(T)[:10])
        axarr[z, 0].plot(strikes, market, label='Market Quotes')
        axarr[z, 0].plot(strikes, model, 'ro', label='Model Prices')
        axarr[z, 0].grid()
        if T is mats[0]:
            axarr[z, 0].set_title('Option Quotes')
        if T is mats[-1]:
            axarr[z, 0].set_xlabel('Strike')
        wi = 2.
        if relative is True:
            axarr[z, 1].bar(strikes - wi / 2,
                           (model - market) / market * 100, width=wi)
        else:
            axarr[z, 1].bar(strikes - wi / 2, model - market, width=wi)
        axarr[z, 1].grid()
        if T is mats[0]:
            axarr[z, 1].set_title('Differences')
        if T is mats[-1]:
            axarr[z, 1].set_xlabel('Strike')
        z += 1


from itertools import product
import numpy as np
from numpy.testing import assert_almost_equal, assert_array_almost_equal

from sklearn import datasets
from sklearn import manifold
from sklearn import neighbors
from sklearn import pipeline
from sklearn import preprocessing
from sklearn.utils.testing import assert_less

eigen_solvers = ['auto', 'dense', 'arpack']
path_methods = ['auto', 'FW', 'D']

def test_isomap_simple_grid():
    # Isomap should preserve distances when all neighbors are used
    N_per_side = 5
    Npts = N_per_side ** 2
    n_neighbors = Npts - 1

    # grid of equidistant points in 2D, n_components = n_dim
    X = np.array(list(product(range(N_per_side), repeat=2)))

    # distances from each point to all others
    G = neighbors.kneighbors_graph(X, n_neighbors,
                                   mode='distance').toarray()

    for eigen_solver in eigen_solvers:
        for path_method in path_methods:
            clf = manifold.Isomap(n_neighbors=n_neighbors, n_components=2,
                                  eigen_solver=eigen_solver,
                                  path_method=path_method)
            clf.fit(X)

            G_iso = neighbors.kneighbors_graph(clf.embedding_,
                                               n_neighbors,
                                               mode='distance').toarray()
            assert_array_almost_equal(G, G_iso)

def test_isomap_reconstruction_error():
    # Same setup as in test_isomap_simple_grid, with an added dimension
    N_per_side = 5
    Npts = N_per_side ** 2
    n_neighbors = Npts - 1

    # grid of equidistant points in 2D, n_components = n_dim
    X = np.array(list(product(range(N_per_side), repeat=2)))

    # add noise in a third dimension
    rng = np.random.RandomState(0)
    noise = 0.1 * rng.randn(Npts, 1)
    X = np.concatenate((X, noise), 1)

    # compute input kernel
    G = neighbors.kneighbors_graph(X, n_neighbors,
                                   mode='distance').toarray()

    centerer = preprocessing.KernelCenterer()
    K = centerer.fit_transform(-0.5 * G ** 2)

    for eigen_solver in eigen_solvers:
        for path_method in path_methods:
            clf = manifold.Isomap(n_neighbors=n_neighbors, n_components=2,
                                  eigen_solver=eigen_solver,
                                  path_method=path_method)
            clf.fit(X)

            # compute output kernel
            G_iso = neighbors.kneighbors_graph(clf.embedding_,
                                               n_neighbors,
                                               mode='distance').toarray()

            K_iso = centerer.fit_transform(-0.5 * G_iso ** 2)

            # make sure error agrees
            reconstruction_error = np.linalg.norm(K - K_iso) / Npts
            assert_almost_equal(reconstruction_error,
                                clf.reconstruction_error())

def test_transform():
    n_samples = 200
    n_components = 10
    noise_scale = 0.01

    # Create S-curve dataset
    X, y = datasets.samples_generator.make_s_curve(n_samples, random_state=0)

    # Compute isomap embedding
    iso = manifold.Isomap(n_components, 2)
    X_iso = iso.fit_transform(X)

    # Re-embed a noisy version of the points
    rng = np.random.RandomState(0)
    noise = noise_scale * rng.randn(*X.shape)
    X_iso2 = iso.transform(X + noise)

    # Make sure the rms error on re-embedding is comparable to noise_scale
    assert_less(np.sqrt(np.mean((X_iso - X_iso2) ** 2)), 2 * noise_scale)

def test_pipeline():
    # check that Isomap works fine as a transformer in a Pipeline
    # only checks that no error is raised.
    # TODO check that it actually does something useful
    X, y = datasets.make_blobs(random_state=0)
    clf = pipeline.Pipeline(
        [('isomap', manifold.Isomap()),
         ('clf', neighbors.KNeighborsClassifier())])
    clf.fit(X, y)
    assert_less(.9, clf.score(X, y))