# -*- coding: utf-8 -*-
"""
Created on Tue Mar 18 11:37:28 2025

@author: Jiawen
"""

'''
this script removes the dark frames from dataset as if they never existed
plist also removed and their frame number replaced by succeeding frames
be careful with pixel value threshold and check removed frames!

'''
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from functions import *
from skimage import io
import plistlib as plib
import glob
from natsort import natsorted
#%%
threshold = 530 #this normally works for XiD (CMOS), check for other cameras.
data_fp = r'F:\05.12.25 piv fish 2\piv 1\Red'
#change this fp for the other zstack
tif_fp = glob.glob(data_fp+'/*.tif')
plist_fp = glob.glob(data_fp+'/*.plist')
plist_fp = natsorted(plist_fp) #added to make sure the laoded plist actually correspond to file paths
tif_fp = natsorted(tif_fp)
total_plist = load_plist(plist_fp)
#%%
dark_framenum=[]
dark_frame=[]
for i in tqdm(range(0, len(tif_fp))):
    sequence = io.imread(tif_fp[i]) #load in images
    for j in range(0, len(sequence)):
        if np.mean(sequence[j]) < threshold:
            #check average pixel intensity with threshold
            #record frame number of frames with lower intensity -> dark frame
            frame_num = total_plist[i].get('frames')[j].get('frame_number')
            print('dark frame found at:' + str(i)+'th tif, '+str(j)+'th image')
            dark_frame.append(sequence[j])
            dark_framenum.append(frame_num)
            


#%% visually check dark frames are actually dark
for frame in dark_frame:
    plt.imshow(frame)
    plt.show()
#%%
#this section is no longer in use, as we no longer delete dark frames
#instead we mark them with a negative frame number in their plists, so they
#do not get involved with frame pair selection.
#modifying raw image data is not a good idea. plists are easier to backup,
#and they can always be restored as we know how much they were changed.
'''
for i in tqdm(range(0, len(tif_fp))):
    #new_plist = total_plist[i].copy()
    sequence = io.imread(tif_fp[i])
    tbd=[]
    for j in range(0, len(sequence)):
        if total_plist[i].get('frames')[j].get('frame_number') in dark_framenum:
            tbd.append(j)
        #deleting dark frames, saving new images
    new_sequence = np.delete(sequence, tbd, axis=0)
    io.imsave(tif_fp[i], new_sequence)
'''
#%%
new_total_plist=[]
frame_num_change = [] #recording how much every frame number should be changed. 
change = 0#after one dark frame, all subseqeunt frame no. should be deducted by 1.
for i in range(0, len(total_plist)): #removing entries from dark frames 
    new_plist = total_plist[i].copy()
    tbd = []
    row_change=[]
    for j in range(0, len(total_plist[i].get('frames'))):
        if total_plist[i].get('frames')[j].get('frame_number') in dark_framenum:
            tbd.append(j)
            change+=1
            row_change.append(change) #added because dark frame plists are no longer removed
        else:
            row_change.append(change)
    frame_num_change.append(row_change)
    #new_plist['frames'] = [item for i, item in enumerate(new_plist['frames']) if i not in tbd]
    for ind in tbd:
        new_plist['frames'][ind]['frame_number'] *= -1 #instead of deleting plists, make frame number negative so they cannot be selected, and also has a chance to restore
    new_total_plist.append(new_plist)

#%% changing frame numbers of all frames after dark frames
for i in range(0, len(new_total_plist)):
    for j in range(0, len(new_total_plist[i].get('frames'))):
        new_total_plist[i]['frames'][j]['frame_number'] -= frame_num_change[i][j]






#%% verifying that there are no pair switching problem
SAD = []
fnum = []
phase = []
#total_plist = load_plist(plist_fp) #loading processed plists, use this if skipped previous sections, e.g. rerun for an old dataset
total_plist = new_total_plist
for i in tqdm(range(0, len(tif_fp))):
    sequence = io.imread(tif_fp[i]) #load in images
    for j in range(0, len(sequence)):
        if j == len(sequence) - 1:
            last_frame = sequence[j]
        if i==0 and j==0:
            pass
        elif i!=0 and j==0:
            frame1 = last_frame
            frame2 = sequence[j]
            SAD.append(np.sum(np.abs(frame2.astype(np.int64)-frame1.astype(np.int64)))) #converting to 64bit to prevent integer overflow 
            fnum.append(total_plist[i].get('frames')[j].get('frame_number'))
            phase.append(total_plist[i].get('frames')[j].get('phase_from_offline_sync_analysis'))
        elif i!=0 and j!=0:
            frame1 = sequence[j-1]
            frame2 = sequence[j]
            SAD.append(np.sum(np.abs(frame2.astype(np.int64)-frame1.astype(np.int64))))
            fnum.append(total_plist[i].get('frames')[j].get('frame_number'))
            phase.append(total_plist[i].get('frames')[j].get('phase_from_offline_sync_analysis'))
#%% separate SAD between frame pairs into even and odd pairs, within a phase range
evenSAD=[]
oddSAD=[]
even_fnum = []
odd_fnum = []
phase_value = 4
phase = [-1 if x is None else x for x in phase]

for i in range(0, len(SAD)):
    if fnum[i]%2 == 0 and fnum[i]>0 and phase[i]>phase_value-0.05 and phase[i]<phase_value+0.05:
        evenSAD.append(SAD[i])
        even_fnum.append(fnum[i])
    elif fnum[i]>0 and phase[i]>phase_value-0.05 and phase[i]<phase_value+0.05:
        oddSAD.append(SAD[i])
        odd_fnum.append(fnum[i])
plt.plot(even_fnum, evenSAD, '.', label='even pairs', markersize=2)
plt.plot(odd_fnum, oddSAD, '.', label='odd pairs', markersize=2)
plt.xlabel('frame pairs')
plt.ylabel('frame pair SAD')
plt.legend()
plt.title('SAD of frame pairs after pair switching fix')
plt.show()

#%% when verified that there are no more pair switching, save new plist files.
for i in tqdm(range(0, len(plist_fp))):
    with open(plist_fp[i], 'wb') as file:
        plib.dump(new_total_plist[i], file)
    file.close()











