#cellpose env
import os
import numpy as np
import pandas as pd
from cellpose import models
import tifffile
from skimage.measure import regionprops_table
import re

# Function to tile images
def tile_images(images, height, width, num_rows = 3, num_cols =4):
    num_images = len(images)
    
    tiled_image = np.zeros((num_rows * height, num_cols * width), dtype=np.uint16)
    
    for i in range(num_rows):
        for j in range(num_cols):
            idx = i * num_cols + j
            if idx < num_images:
                tiled_image[i * height:(i + 1) * height, j * width:(j + 1) * width] = images[idx]
    
    return tiled_image

#set path to folder for max projected images
im_path  = "[path_to_max_projected_images]"

#set save_path
save_path = "[path_to_save_results]"

#set path to your cellpose model - if using default model this can be changed
model_path = "cyto3"


#set the number of rows and columns for each well to tile
rows = 3
cols = 4

#get the list of files in the projected image folder
im_files = os.listdir(im_path)

#identify files for image segmentation
C2_files = [file for file in im_files if file.endswith("C2.tif") and not file.endswith("SUMC2.tif")]

#generate dataframe of raw images
im_df= pd.DataFrame()

im_df['C2_file'] = C2_files
im_df['path'] = im_path
im_df['well'] = im_df['C2_file'].apply(lambda x: int(x.split('W')[1].split('F')[0]))
im_df['field'] = im_df['C2_file'].apply(lambda x: int(x.split('F')[1].split('T')[0]))
im_df['time'] = im_df['C2_file'].apply(lambda x: int(x.split('T')[1].split('Z')[0]))
#im_df['z'] = im_df['C2_file'].str.extract('Z(\d+)').astype(int)
im_df['C2_file']
im_df['C2_file'] = im_df.apply(lambda row: os.path.join(row['path'], row['C2_file']), axis=1)
im_df['C3_file'] = im_df['C2_file'].str.replace('C2.tif', 'C3.tif')
im_df['C1_file'] = im_df['C2_file'].str.replace('C2.tif', 'C1.tif')
im_df['C4_file'] = im_df['C2_file'].str.replace('C2.tif', 'C4.tif')

im_df = im_df.drop('path', axis=1)

#group by well and time points
os.makedirs("./tiled_images", exist_ok= True)
os.makedirs("./quant", exist_ok= True)

grouped = im_df.sort_values(by=['well','time'])
grouped = grouped.groupby(['well', 'time'])


#tile and segment images and save
for name, group_df in grouped:
    well, time = name
    print("W" + str(f'{well}').zfill(3) + "T" + str(f'{time}').zfill(3))
    if os.path.exists(save_path + "/tiled_images/"+ "W" + str(f'{well}').zfill(3) + "T" + str(f'{time}').zfill(3) + "_tiled_stack.tif"):
        print("Exists...")
        continue
 
    images = []
    for idx, row in group_df.iterrows():
        well_num = row['well']
        #if well_num ==14:
        #    continue
        tiff_path = row['C1_file']
        tiff_image = tifffile.imread(tiff_path)
        images.append(tiff_image)
    #if well_num ==14:
    #    continue

    height, width = images[0].shape
    
    # Tile images
    tiled_image = tile_images(images, height, width, num_rows = rows, num_cols = cols)

    #plt.imshow(tiled_image)
    #plt.show()
    im_stack = np.zeros((5, tiled_image.shape[0], tiled_image.shape[1]), dtype = tiled_image.dtype)
    im_stack[0] = tiled_image

    images = []
    for idx, row in group_df.iterrows():

        tiff_path = row['C2_file']
        tiff_image = tifffile.imread(tiff_path)
        images.append(tiff_image)

    height, width = images[0].shape
    
    # Tile images
    tiled_image = tile_images(images, height, width, num_rows = rows, num_cols = cols)

    #plt.imshow(tiled_image)
    #plt.show()
    im_stack[1] = tiled_image

    images = []
    for idx, row in group_df.iterrows():

        tiff_path = row['C3_file']
        tiff_image = tifffile.imread(tiff_path)
        images.append(tiff_image)

    height, width = images[0].shape
    
    # Tile images
    tiled_image = tile_images(images, height, width, num_rows = rows, num_cols = cols)

    #plt.imshow(tiled_image)
    #plt.show()
    im_stack[2] = tiled_image

    images = []
    for idx, row in group_df.iterrows():

        tiff_path = row['C4_file']
        tiff_image = tifffile.imread(tiff_path)
        images.append(tiff_image)

    height, width = images[0].shape
    
    # Tile images
    tiled_image = tile_images(images, height, width, num_rows = rows, num_cols = cols)

    im_stack[3] = tiled_image


    print("Segmenting")

    #specify local model

    model = models.CellposeModel(gpu=True, model_type= model_path)
    im_seg = im_stack[0:3,:,:]       

    masks, flows, styles = model.eval(im_seg, diameter=150, channels = [3,2],  do_3D = False, batch_size=32)

    im_stack[4] = masks

    #save segmented images - can comment out if not needed
    tifffile.imwrite(save_path + "/tiled_images/"+ "W" + str(f'{well}').zfill(3) + "T" + str(f'{time}').zfill(3) + "_tiled_stack.tif", im_stack, bigtiff=True, compression = "zlib")

    #quantify cell morphology and image intensity
    quant = pd.DataFrame(regionprops_table(label_image = masks, intensity_image = im_stack[3,:,:],  properties=['label', 'area', 'axis_major_length', 'axis_minor_length', 'eccentricity', 'centroid', 'mean_intensity']))
    
    #save
    quant.to_csv(save_path + "/quant/"+ "W" + str(f'{well}').zfill(3) + "T" + str(f'{time}').zfill(3) + "_quant.csv", index=False)