import h5py
import csv
import numpy
#------------------------------------------ INPUT PARAMETERS ---------------------------------------
# Number and name of files for time-resolved results

nb_files = 26
filenumber = [None] * nb_files
filenumber[0] = '535428'
filenumber[1] = '535429'
filenumber[2] = '535432'
filenumber[3] = '535433'
filenumber[4] = '535438'
filenumber[5] = '535439'
filenumber[6] = '535440'
filenumber[7] = '535443'
filenumber[8] = '535444'
filenumber[9] = '535445'
filenumber[10] = '535448'
filenumber[11] = '535449'
filenumber[12] = '535450'
filenumber[13] = '535453'
filenumber[14] = '535454'
filenumber[15] = '535455'
filenumber[16] = '535458'
filenumber[17] = '535459'
filenumber[18] = '535460'
filenumber[19] = '535464'
filenumber[20] = '535465'
filenumber[21] = '535469'
filenumber[22] = '535470'
filenumber[23] = '535473'
filenumber[24] = '535474'
filenumber[25] = '535475'

index_file = 0
saxs_sub_time = []
waxs_sub_time = []
saxs_div_time = []
waxs_div_time = []

for index_file in range(nb_files) :
    print(filenumber[index_file])
    #------------------------------------------ INIT ---------------------------------------------------
    nb_goodframe = 0
    saxs_total_int = 0
    waxs_total_int = 0
    saxs_avg_int = 0
    waxs_avg_int = 0
    saxs=[]
    saxs_sub=[]
    waxs_sub=[]
    saxs_div=[]
    waxs_div=[]
    waxs=[]
    R = [None] * 1000
    R2 = [None] * 1000
    
    nb_1_frames = 0

    saxs_region_1_tot = 0
    waxs_region_1_tot = 0
    saxs_region_1_avg = 0
    waxs_region_1_avg = 0
    
    saxs_region_1 = []
    waxs_region_1 = []
    
    saxs_max = []
    saxs_min= []
    waxs_min =[]
    ind_max = 0
    ind_min = 0
    
    int_disc_1 = 0
    #----------------------------------------- FILES OPENING AND DATA READING -------------------------
    #SAXS-WAXS filenames, outputs from the beamline software
    saxs_filename = 'i22-'+filenumber[index_file]+'_saxs_Transmission_IvsQ_processed.nxs'
    waxs_filename = 'i22-'+filenumber[index_file]+'_waxs_Transmission_IvsQ_processed.nxs'
    #open files
    f = h5py.File(saxs_filename, 'r')
    g = h5py.File(waxs_filename, 'r')
    #definition of datasets as arrays
    saxs_q_range = f['processed/result/q'][:]  #q range values
    waxs_q_range = g['processed/result/q'][:]
    
    saxs_scattered_int = f['processed/result/data'][0,:,:]  #intensity vs frame number and q range
    waxs_scattered_int = g['processed/result/data'][0,:,:]
        
    #-------------------------------------------------------------------------------------------------
    #----------------- FRAMES SORTING BASED on PEAK INTENSITIES COMPARISONAND AVERAGING --------------
    for index_frame in range(1000):  #loop on all the 1000 frames in one dataset
        R[index_frame] = waxs_scattered_int[index_frame,357] / waxs_scattered_int[index_frame,250]
        R2[index_frame] = waxs_scattered_int[index_frame,360] / waxs_scattered_int[index_frame,290]
        if R[index_frame] > 1 and R2[index_frame] > 1:
            saxs_total_int = saxs_total_int + saxs_scattered_int[index_frame,:] #sum of goof frames     
            waxs_total_int = waxs_total_int + waxs_scattered_int[index_frame,:]          
            nb_goodframe = nb_goodframe + 1   #incrementing good frames number     
            saxs.append(saxs_scattered_int[index_frame,:]) #new arrays with good frames only
            waxs.append(waxs_scattered_int[index_frame,:])
    
    saxs_avg_int = saxs_total_int / nb_goodframe #averaging the sum of good frames
    waxs_avg_int = waxs_total_int / nb_goodframe
    #close files
    f.close()
    g.close()
    #-------------------------------- WRITING NEW DATASETs AND AVERAGE DATA -------------------------
    #write datasets without oil frames in a new .h5 file
    name_file_nooil = filenumber[index_file]+'_nooil_peak.h5'
    f = h5py.File(name_file_nooil, 'w')
    f.create_dataset('saxs', data=saxs)
    f.create_dataset('saxs_q', data=saxs_q_range)
    f.create_dataset('waxs_q', data=waxs_q_range)
    f.create_dataset('waxs', data=waxs)
    f.create_dataset('avg_saxs', data=saxs_avg_int)
    f.create_dataset('avg_waxs', data=waxs_avg_int)
    f.create_dataset('R', data=R)
    f.create_dataset('R2', data=R2)
    
    sorted_saxs = f['/saxs'][:,:]  #intensity as function of frame number and q range
    sorted_waxs = f['/waxs'][:,:] 
    
    #------------------------------------------------------------------------------
    #----------------- SORTING NANOPARTICLES/SOLUTION FRAMES ----------------------
    #Picking max(nanoparticle) and min(solution) frames
    saxs_min = sorted_saxs[0,:]
    saxs_max = sorted_saxs[0,:]
    
    for index_frame in range(nb_goodframe):        
        if sorted_saxs[index_frame,400]>saxs_max[400]:
            saxs_max = sorted_saxs[index_frame,:]
            ind_max = index_frame
                    
    for index_frame in range(nb_goodframe):    
        if sorted_saxs[index_frame,400]<saxs_min[400]:
            saxs_min = sorted_saxs[index_frame,:]
            ind_min = index_frame
    
    saxs_max = sorted_saxs[ind_max,:]
    saxs_min = sorted_saxs[ind_min,:]
    waxs_min = sorted_waxs[ind_min,:]
	waxs_max = sorted_waxs[ind_max,:]
        
    #Averaging Low SAXS Intensities (Solution)
    saxs_disc_1 = numpy.exp((7*numpy.log(saxs_min)+numpy.log(saxs_max))/8)
    f.create_dataset('saxs_disc_1', data=saxs_disc_1)
    
    #Picking intermediate value for intensity comparison
    int_disc_1 = saxs_disc_1[400]
    
    #Sorting low intensity frames
    for index_frame in range(nb_goodframe):
         
        if sorted_saxs[index_frame,400] < int_disc_1: #Region 1 Low intensity 
            nb_1_frames += 1
            saxs_region_1.append(sorted_saxs[index_frame,:])
            waxs_region_1.append(sorted_waxs[index_frame,:])
            
            saxs_region_1_tot = saxs_region_1_tot + sorted_saxs[index_frame,:]
            waxs_region_1_tot = waxs_region_1_tot + sorted_waxs[index_frame,:]

            
    prop_frames = nb_1_frames/nb_goodframe *100
                         
    #Averaging low intensities
    if nb_1_frames != 0 :
        saxs_region_1_avg = saxs_region_1_tot / nb_1_frames
        waxs_region_1_avg = waxs_region_1_tot / nb_1_frames

        
    #------------------SUBSTRACTED/DIVIDED DATA---------------------
    for index_frame in range(nb_goodframe):
            #Substracting min frame    
        saxs_sub = sorted_saxs[index_frame,:] - saxs_region_1_avg[:]
        saxs_sub_time.append(saxs_sub)
        waxs_sub = sorted_waxs[index_frame,:] - waxs_region_1_avg[:]
        waxs_sub_time.append(waxs_sub)
        saxs_div = sorted_saxs[index_frame,:] / saxs_region_1_avg[:]
        saxs_div_time.append(saxs_div)
        waxs_div = sorted_waxs[index_frame,:] / waxs_region_1_avg[:]
        waxs_div_time.append(waxs_div)
        
    #Creating new datasets    
    f.create_dataset('prop_frames', data=prop_frames)
    f.create_dataset('saxs_region_1_avg', data=saxs_region_1_avg)
    f.create_dataset('waxs_region_1_avg', data=waxs_region_1_avg)
    f.create_dataset('saxs_max', data=saxs_max)
    f.create_dataset('saxs_min', data=saxs_min)
    f.close()
     
    
    #write new .dat files with the average intensity as a function of q for SAXS, WAXS
    ex_saxs_filename = filenumber+'_saxs_average.dat'
    ex_waxs_filename = filenumber+'_waxs_average.dat'
    #saxs
    with open(ex_saxs_filename, 'w') as f:
        writer = csv.writer(f, delimiter='\t')
        writer.writerows(zip(saxs_q_range,saxs_avg_int))
    f.close()
    #waxs
    with open(ex_waxs_filename, 'w') as f:
        writer = csv.writer(f, delimiter='\t')
        writer.writerows(zip(waxs_q_range,waxs_avg_int))
    f.close()
    #------------------------------------------------------------------------------------------------

#-------------------------TIME RESOLVED DATA-------------------------
f = h5py.File('time_resolved_sub.h5', 'w')
f.create_dataset('saxs_sub', data=saxs_sub_time)
f.create_dataset('waxs_sub', data=waxs_sub_time)
f.create_dataset('saxs_div', data=saxs_div_time)
f.create_dataset('waxs_div', data=waxs_div_time)
f.create_dataset('saxs_q', data=saxs_q_range)
f.create_dataset('waxs_q', data=waxs_q_range)
f.close()