"""
https://elog-gfa.psi.ch/SwissFEL+commissioning/19324


5.06 mm /sf/data/measurements/2021/04/25/2021_04_25-17_16_30_Lasing_False_SARBD02-DSCR050.h5

5.04 mm /sf/data/measurements/2021/04/25/2021_04_25-17_33_31_Lasing_False_SARBD02-DSCR050.h5

5.02 mm /sf/data/measurements/2021/04/25/2021_04_25-17_36_28_Lasing_False_SARBD02-DSCR050.h5

5.00 mm /sf/data/measurements/2021/04/25/2021_04_25-17_37_20_Lasing_False_SARBD02-DSCR050.h5

-4.4 mm /sf/data/measurements/2021/04/25/2021_04_25-17_40_15_Lasing_False_SARBD02-DSCR050.h5

-4.42 mm /sf/data/measurements/2021/04/25/2021_04_25-17_41_26_Lasing_False_SARBD02-DSCR050.h5

-4.38 mm /sf/data/measurements/2021/04/25/2021_04_25-17_42_51_Lasing_False_SARBD02-DSCR050.h5

-4.36 mm /sf/data/measurements/2021/04/25/2021_04_25-17_44_59_Lasing_False_SARBD02-DSCR050.h5

Lasing ON 510 uJ

-4.42 mm /sf/data/measurements/2021/04/25/2021_04_25-17_49_14_Lasing_True_SARBD02-DSCR050.h5

-4.40 mm /sf/data/measurements/2021/04/25/2021_04_25-17_50_14_Lasing_True_SARBD02-DSCR050.h5

-4.38 mm /sf/data/measurements/2021/04/25/2021_04_25-17_51_35_Lasing_True_SARBD02-DSCR050.h5

-4.36 mm /sf/data/measurements/2021/04/25/2021_04_25-17_52_38_Lasing_True_SARBD02-DSCR050.h5

5.00 mm /sf/data/measurements/2021/04/25/2021_04_25-17_54_39_Lasing_True_SARBD02-DSCR050.h5

5.02 mm /sf/data/measurements/2021/04/25/2021_04_25-17_55_43_Lasing_True_SARBD02-DSCR050.h5

5.04 mm /sf/data/measurements/2021/04/25/2021_04_25-17_56_38_Lasing_True_SARBD02-DSCR050.h5

5.06 mm /sf/data/measurements/2021/04/25/2021_04_25-17_57_34_Lasing_True_SARBD02-DSCR050.h5

Changing S30CB15 first passive streaker (10 mm gap) 4.72 offset

SARUN18 offset, S30CB15 offset, filename

5.06 mm 4.41 mm /sf/data/measurements/2021/04/25/2021_04_25-18_16_36_Lasing_True_SARBD02-DSCR050.h5

5.06 mm 4.61 mm /sf/data/measurements/2021/04/25/2021_04_25-18_17_56_Lasing_True_SARBD02-DSCR050.h5

5.06 mm 4.78 mm /sf/data/measurements/2021/04/25/2021_04_25-18_19_02_Lasing_True_SARBD02-DSCR050.h5

5.06 mm 4.78 mm /sf/data/measurements/2021/04/25/2021_04_25-18_19_58_Lasing_False_SARBD02-DSCR050.h5

5.06 mm 4.61 mm /sf/data/measurements/2021/04/25/2021_04_25-18_20_48_Lasing_False_SARBD02-DSCR050.h5

5.06 mm 4.41 mm /sf/data/measurements/2021/04/25/2021_04_25-18_21_41_Lasing_False_SARBD02-DSCR050.h5


2 color sext

2nd sext in bc1 originally at 0.105 A

3 A k=4 /sf/data/measurements/2021/04/25/2021_04_25-18_55_46_Lasing_True_SARBD02-DSCR050.h5

4 A k=2 /sf/data/measurements/2021/04/25/2021_04_25-18_58_36_Lasing_True_SARBD02-DSCR050.h5

4 A k=3 /sf/data/measurements/2021/04/25/2021_04_25-18_59_54_Lasing_True_SARBD02-DSCR050.h5 /sf/data/measurements/2021/04/25/2021_04_25-19_01_13_Lasing_False_SARBD02-DSCR050.h5

6 A k=2 /sf/data/measurements/2021/04/25/2021_04_25-19_03_16_Lasing_True_SARBD02-DSCR050.h5 /sf/data/measurements/2021/04/25/2021_04_25-19_04_24_Lasing_False_SARBD02-DSCR050.h5

6 D k=2.5 /sf/data/measurements/2021/04/25/2021_04_25-19_05_39_Lasing_True_SARBD02-DSCR050.h5 /sf/data/measurements/2021/04/25/2021_04_25-19_06_55_Lasing_False_SARBD02-DSCR050.h5

6 A k = 0.5 /sf/data/measurements/2021/04/25/2021_04_25-19_10_07_Lasing_True_SARBD02-DSCR050.h5



above all k were 1

4 A k=2 /sf/data/measurements/2021/04/25/2021_04_25-19_20_58_Lasing_True_SARBD02-DSCR050.h5 pulseEne = 215 uJ

4 A k=3 /sf/data/measurements/2021/04/25/2021_04_25-19_23_46_Lasing_True_SARBD02-DSCR050.h5 pulseEne = 100 uJ

4 A k=1 /sf/data/measurements/2021/04/25/2021_04_25-19_30_22_Lasing_True_SARBD02-DSCR050.h5 pulseEne = 360 uJ

6 A k=1 pulseEne=100 uJ /sf/data/measurements/2021/04/25/2021_04_25-19_40_02_Lasing_True_SARBD02-DSCR050.h5 /sf/data/measurements/2021/04/25/2021_04_25-19_41_05_Lasing_False_SARBD02-DSCR050.h5

6 A k=0.5 pulseEne= 220 uJ /sf/data/measurements/2021/04/25/2021_04_25-19_46_03_Lasing_True_SARBD02-DSCR050.h5

6 A k=0.75 pulseEne=155 uJ /sf/data/measurements/2021/04/25/2021_04_25-19_49_04_Lasing_True_SARBD02-DSCR050.h5

6 A k=1.5 pulseEne=50 uJ /sf/data/measurements/2021/04/25/2021_04_25-19_52_08_Lasing_True_SARBD02-DSCR050.h5

6 A k=1.25 pulseEne=71 uJ /sf/data/measurements/2021/04/25/2021_04_25-19_55_02_Lasing_True_SARBD02-DSCR050.h5
"""


import os
import numpy as np
import matplotlib.pyplot as plt
import copy
from socket import gethostname

from h5_storage import loadH5Recursive
import tracking
import image_and_profile as iap
import elegant_matrix
import misc2 as misc
import myplotstyle as ms

elegant_matrix.set_tmp_dir('~/tmp_elegant/')

plt.close('all')


hostname = gethostname()
if hostname == 'desktop':
    data_dir = '/storage/data_2021-04-25/'
elif hostname == 'pc11292.psi.ch':
    data_dir = '/sf/data/measurements/2021/04/25/'
elif hostname == 'pubuntu':
    data_dir = '/storage/data_2021-04-25/'


offset_file_dict = {
        'OFF': {
            5.06: '/sf/data/measurements/2021/04/25/2021_04_25-17_16_30_Lasing_False_SARBD02-DSCR050.h5',
            5.04: '/sf/data/measurements/2021/04/25/2021_04_25-17_33_31_Lasing_False_SARBD02-DSCR050.h5',
            5.02: '/sf/data/measurements/2021/04/25/2021_04_25-17_36_28_Lasing_False_SARBD02-DSCR050.h5',
            5.00: '/sf/data/measurements/2021/04/25/2021_04_25-17_37_20_Lasing_False_SARBD02-DSCR050.h5',
            -4.4: '/sf/data/measurements/2021/04/25/2021_04_25-17_40_15_Lasing_False_SARBD02-DSCR050.h5',
            -4.42: '/sf/data/measurements/2021/04/25/2021_04_25-17_41_26_Lasing_False_SARBD02-DSCR050.h5',
            -4.38: '/sf/data/measurements/2021/04/25/2021_04_25-17_42_51_Lasing_False_SARBD02-DSCR050.h5',
            -4.36: '/sf/data/measurements/2021/04/25/2021_04_25-17_44_59_Lasing_False_SARBD02-DSCR050.h5',
            },
        'ON': {
            -4.42: '/sf/data/measurements/2021/04/25/2021_04_25-17_49_14_Lasing_True_SARBD02-DSCR050.h5',
            -4.40: '/sf/data/measurements/2021/04/25/2021_04_25-17_50_14_Lasing_True_SARBD02-DSCR050.h5',
            -4.38: '/sf/data/measurements/2021/04/25/2021_04_25-17_51_35_Lasing_True_SARBD02-DSCR050.h5',
            -4.36: '/sf/data/measurements/2021/04/25/2021_04_25-17_52_38_Lasing_True_SARBD02-DSCR050.h5',
            5.00: '/sf/data/measurements/2021/04/25/2021_04_25-17_54_39_Lasing_True_SARBD02-DSCR050.h5',
            5.02: '/sf/data/measurements/2021/04/25/2021_04_25-17_55_43_Lasing_True_SARBD02-DSCR050.h5',
            5.04: '/sf/data/measurements/2021/04/25/2021_04_25-17_56_38_Lasing_True_SARBD02-DSCR050.h5',
            5.06: '/sf/data/measurements/2021/04/25/2021_04_25-17_57_34_Lasing_True_SARBD02-DSCR050.h5',
            },
        }

all_offsets = np.array(sorted(offset_file_dict['ON'].keys()))

analysis_file = '/storage/tmp_reconstruction/2021_04_29-16_06_46_PassiveReconstruction.h5'
analysis_dict = loadH5Recursive(analysis_file)
streaker_center = analysis_dict['input']['streaker_means'][1]
screen_x0 = analysis_dict['input']['screen_x0']
tracker_kwargs = analysis_dict['input']['tracker_kwargs']

tracker_kwargs['quad_wake'] = True
tracker_kwargs['override_quad_beamsize'] = False


gauss_kwargs = analysis_dict['input']['gaussian_reconstruction']
tracker = tracking.Tracker(**tracker_kwargs)

ms.figure('Current profiles')
subplot = ms.subplot_factory(2,2)
sp_ctr = 1

sp_profile0 = subplot(sp_ctr, title='Reconstructed profiles', xlabel='t (fs)', ylabel='I (kA)')
sp_ctr += 1

sp_screen0 = subplot(sp_ctr, title='Screens', xlabel='x (mm)', ylabel='Intensity (arb. units)')
sp_ctr += 1

sp_profile1 = subplot(sp_ctr, title='Reconstructed profiles', xlabel='t (fs)', ylabel='I (kA)')
sp_ctr += 1

sp_screen1 = subplot(sp_ctr, title='Screens', xlabel='x (mm)', ylabel='Intensity (arb. units)')
sp_ctr += 1

for streaker_offset_mm in all_offsets:

    if streaker_offset_mm > 0:
        sp_screen, sp_profile = sp_screen0, sp_profile0
    if streaker_offset_mm < 0:
        sp_screen, sp_profile = sp_screen1, sp_profile1

    streaker_offset = streaker_offset_mm*1e-3

    file_off0 = offset_file_dict['OFF'][streaker_offset_mm]
    file_off = data_dir+os.path.basename(file_off0)
    dict_off = loadH5Recursive(file_off)['pyscan_result']
    all_proj = dict_off['image'].astype(float).sum(axis=-2)
    x_axis = dict_off['x_axis'].astype(float)
    median_index = misc.get_median(all_proj, output='index')
    median_proj = all_proj[median_index]
    meas_screen = misc.proj_to_screen(median_proj, x_axis, True, screen_x0)


    gauss_kwargs2 = copy.deepcopy(gauss_kwargs)
    gauss_kwargs2['beam_offsets'] = [0, -(streaker_offset - streaker_center)]
    del gauss_kwargs2['meas_screen_intensity']
    del gauss_kwargs2['meas_screen_x']
    gauss_kwargs2['meas_screen'] = meas_screen
    gauss_kwargs2['sig_t_range'] = np.arange(20, 50.01, 5)*1e-15
    gauss_dict = tracker.find_best_gauss(**gauss_kwargs2)

    profile = gauss_dict['reconstructed_profile']
    screen = gauss_dict['reconstructed_screen']

    _label='%.2f mm %i fs' % (streaker_offset_mm, profile.rms()*1e15)
    profile.plot_standard(sp_profile, label=_label, center='Gauss')

    _label = '%.2f mm' % streaker_offset_mm
    color = meas_screen.plot_standard(sp_screen, label=_label)[0].get_color()

    screen.plot_standard(sp_screen, color=color, ls='--')

    file_on0 = offset_file_dict['ON'][streaker_offset_mm]
    file_on = data_dir+os.path.basename(file_on0)
    dict_on = loadH5Recursive(file_on)
    all_proj = dict_on['image'].astype(float).sum(axis=-2)
    x_axis = dict_on['x_axis'].astype(float)
    median_index = misc.get_median(all_proj, output='index')
    median_image = misc.im





for sp_ in sp_profile0, sp_profile1:
    sp_.set_xlim(-100, 100)


for sp_ in sp_profile0, sp_profile1, sp_screen0, sp_screen1:
    sp_.legend()

plt.show()

