import sys
import os
import numpy as np
import glob
import cv2
import cvtools

if __name__=="__main__":

    if len(sys.argv) != 2:
        print('Please supply a directory')
        exit()

    total_focus_files = glob.glob(os.path.join(sys.argv[1],"*_TotalFocus.png"))

    scales = ['01x', '02x', '04x', '08x', '016x', '032x', '064x', '0128x']

    # process background for each set
    scales = ['01x','02x','04x','08x','016x','032x','064x','0128x']
    for scale in scales:
        group = []
        skip_group = False
        for img_file in total_focus_files:
            if scale in img_file:
                group.append(img_file)

        all_imgs = []

        for img_file in group:
            gray = cvtools.import_image(os.path.dirname(img_file), os.path.basename(img_file), raw=True, color=False)
            gray = cvtools.convert_to_8bit(gray)
            all_imgs.append(gray)

        all_imgs = np.array(all_imgs)

        # median and subtract
        med_img = np.median(all_imgs,0)
        for ind,img_file in enumerate(group):
            bg_sub = all_imgs[ind].astype('float') - med_img
            bg_sub_8bit = bg_sub - np.min(bg_sub)
            bg_sub_8bit = (2**16)*(bg_sub/np.max(bg_sub))
            print('saving: '+ img_file+'_bg_sub.tif')
            cv2.imwrite(img_file+'_bg_sub.tif',bg_sub)

    proc_results = []

    bg_sub_files = glob.glob(os.path.join(sys.argv[1], "*bg_sub*.tif"))

    for scale in scales:
        group = []
        skip_group = False
        for img_file in total_focus_files:
            if scale in img_file:
                group.append(img_file)
        mean_particles = 0
        for ind,img_file in enumerate(group):
            show_proc = False
            if ind == 1:
                show_proc = True
            mean_particles += cvtools.find_particles(img_file, save_to_disk=False, show=show_proc)
        print(mean_particles/len(group))