##  examine the effect of acquire_time and variation of dark current

# fix total exposure time and changing acquire time to observe the
# quality of PDF.

import os
import numpy as np
import matplotlib.pyplot as plt
import tifffile as tif


round_num = input('Which round is this test?  ')


acq_time_list = [0.1, 0.5, 1., 3., 5., 10.] # time per frame
total_exp_time = 300  # adjust based on sample
glbl.dk_window = 0.01 # make sure everytime collect a dark
ex = Experiment('0630test', bt)
sp_str = 'ct_'+ str(total_exp_time)
ScanPlan(sp_str)


for num in acq_time_list:
    print('collecting {} over list {}'.format(num, acq_time_list))
    glbl.frame_acq_time = num # change aquire time first
    sample_str = 'acq_time_'+str(num) # make new sample, so tiff name updates
    Sample(sample_str, ex)
    prun(sample_str, sp_str)
    Tim_save_tiff(db[-1])


# plot and save dark current intensity
dark_files = [fn for fn in os.listdir(glbl.tiff_base) if
              fn.startswith('dark')]
dark_img_container = []
dark_img_int_container = []
for el in dark_files:
    dark_img = tif.imread(os.path.join(glbl.tiff_base, el))
    dark_img_container.append(dark_img)
    dark_img_int_container.append(np.sum(np.sum(dark_img)))


np.save('dark_img_round={}'.format(round_num), dark_img_container)
if os.path.isfile('dark_img_round={}'.format(round_num)):
    print("dark_img_round={} has been saved in current dir"
          .format(round_num))

np.save('dark_int_round={}'.format(round_num), dark_img_int_container)
if os.path.isfile('dark_int_round={}'.format(round_num)):
    print("dark_int_round={} has been saved in current dir"
          .format(round_num))

fig = plt.figure()
plt.plot(acq_time_list, dark_img_int_container)
plt.show()


# load raw files just in case
raw_files = [fn for fn in os.listdir(glbl.tiff_base) if
              fn.startswith('raw')]
raw_img_container = []
for el in raw_files:
    raw_img = tif.imread(os.path.join(glbl.tiff_base, el))
    raw_img_container.append(raw_img)

np.save('raw_img_round={}'.format(round_num), raw_img_container)
if os.path.isfile('raw_img_round={}'.format(round_num)):
    print("raw_img_round={} has been saved in current dir"
          .format(round_num))


# refinement....


# starting the second round of this loop by doing ``run -i test.py``
# again and compare data quality / dark current value to the first round
# Also see if we can apply dark with the same acquire time but from different
# rounds.


###### function used ######
W_DIR = glbl.tiff_base
_fname_field = ['sa_name','sp_name']

def _feature_gen(header):
    ''' generate a human readable file name.

    file name is generated by metadata information in header
    '''
    uid = header.start.uid[:6]
    feature_list = []

    field = header['start']
    for key in _fname_field:

        # get special label
        try:
            if header.start['xp_isdark']:
                feature_list.append('dark')
        except KeyError:
            pass

        try:
            el = field[key]
            # truncate string length
            if len(el)>12:
                value = el[:12]
            else:
                value = el
            # clear space
            feature = [ ch for ch in list(el) if ch!=' ']
            feature_list.append(''.join(feature))
        except KeyError:
            pass # protection to allow missing required fields. This should not happen
    feature_list.append(uid)
    f_name = "_".join(feature_list)
    return f_name

def _timestampstr(timestamp):
    ''' convert timestamp to strftime formate '''
    timestring = datetime.datetime.fromtimestamp(float(timestamp)).strftime('%Y%m%d-%H%M')
    return timestring


def Tim_save_tiff(headers, dark_subtraction=True, *, max_count=None):
    ''' save images obtained from dataBroker as tiff format files.

    Parameters
    ----------
    headers : list
        a list of header objects obtained from a query to dataBroker

    dark_subtraction : bool, optional
        Default is True, which allows dark/background subtraction to 
        be done before saving each image. If header doesn't contain
        necessary information to perform dark subtraction, uncorrected
        image will be saved.

    max_count : int, optional
        The maximum number of events to process per-run.  This can be
        useful to 'preview' an export or if there are corrupted files
        in the data stream (ex from the IOC crashing during data acquisition).
    '''

    F_EXTEN = '.tif' # request from beamline scientist. No difference actually.
    e = '''Can not find a proper dark image applied to this header. 
    Files will be saved but not no dark subtraction will be applied'''
    is_dark_subtracted = False # Flip it only if subtraction is successfully done
    
    # prepare header
    if type(list(headers)[1]) == str:
        header_list = list()
        header_list.append(headers)
    else:
        header_list = headers

    for header in header_list:
        print('Saving your image(s) now....')
        # information at header level
        img_field = _identify_image_field(header)
        dark_img = None
        if 'sc_dk_field_uid' not in header.start:
            warnings.warn("Requested to do dark correction, but header does "
                          "not contain a 'dk_field_uid' entry.  "
                          "Disabling dark subtraction.")
            dark_subtraction = False

        if dark_subtraction:
            dark_uid_appended = header.start['sc_dk_field_uid']
            try:
                # bluesky only looks for uid it defines
                dark_search = {'group': 'XPD',
                               'sc_dark_uid': dark_uid_appended} # the one we need to look up data

                dark_header = db(**dark_search)
                dark_img = np.asarray(get_images(dark_header,
                                                 img_field)).squeeze()
            except ValueError:
                print(e)  # protection. Should not happen
                warnings.warn("Requested to do dark correction, but "
                              "extracting the dark image failed.  Proceeding "
                              "without correction.")
        for ev in get_events(header, fill=True):
            img = ev['data'][img_field]
            ind = ev['seq_num']
            f_name = _feature_gen(header)
            # time when triggering area detector
            event_timestamp = ev['timestamps']['pe1_image']

            f_name = '_'.join([f_name, _timestampstr(event_timestamp)])
            # dark subtration logic
            if dark_img is not None:
                img -= dark_img
                # add prefix if subtracted
                f_name = 'sub_' + f_name

            # complete file name
            if 'temperature' in ev['data']:
                f_name = f_name + '_' + str(ev['data']['temperature']) + 'K'
            # index is still needed as we don't want timestamp in file
            # name down to seconds
            combind_f_name = '{}_{:05d}{}'.format(f_name, ind, F_EXTEN)
            w_name = os.path.join(W_DIR, combind_f_name)
            dark_w_name = os.path.join(W_DIR, 'dark'+combind_f_name)
            raw_w_name = os.path.join(W_DIR, 'raw'+combind_f_name)
            tif.imsave(w_name, img) # subtracted
            tif.imsave(dark_w_name, dark_img) # dark image
            tif.imsave(raw_w_name, dark_img + img) # raw image
            if os.path.isfile(w_name):
                print('subtracted image "%s" has been saved at "%s"' %
                      (combind_f_name, W_DIR))
            else:
                print('Sorry, something went wrong with your tif saving')
                return
            if os.path.isfile(dark_w_name):
                print('dark image "%s" has been saved at "%s"' %
                      (os.path.basename(dark_w_name), W_DIR))
            if os.path.isfile(raw_w_name):
                print('raw image "%s" has been saved at "%s"' %
                      (os.path.basename(raw_w_name), W_DIR))
            if max_count is not None and ind >= max_count:
                # break the loop if max_count reached, move to next header
                break

    print('||********Saving process FINISHED********||')
