#!/usr/bin/env python
# coding: utf-8

# In[1]:


#doen't work on older versions of cellpose. v0.7.2 doesn' run on gpu
import numpy as np
import time, os, sys
import scipy.io
#import os.path
from os import path

#import mxnet as mx
#import matplotlib.pyplot as plt
import glob
import sys

from skimage import io
from skimage import filters

from cellpose import models

# model_type='cyto' or model_type='nuclei'
model = models.Cellpose(gpu=True, model_type='cyto')


filelist = glob.glob('processed/MAX*/aligned/alignedn2vhyb01.tif')
filelist1 = glob.glob('processed/MAX*/aligned/alignedfixedn2vgeneseq01.tif')
folderlist=glob.glob('processed/MAX*/aligned/')
#len(filelist)


# In[2]:



for i in range(len(filelist)): 
    if i%10==0:
        print(i)
    if not path.exists(folderlist[i]+'cellmask_fake.mat'):
        imgs=[]
        
        currfile=filelist[i]
        currfile1=filelist1[i]
        #print(currfile)
        hyb=io.ImageCollection(currfile,conserve_memory=False)
                    #hyb_sum=hyb[0]+hyb[1]+hyb[2]+hyb[3]
                    
        bc=io.ImageCollection(currfile1,conserve_memory=False)
                    #bc_sum=bc[0]+bc[1]+bc[2]+bc[3]
                    #cyto_sum=hyb_sum+bc_sum
        cyto_sum=sum(hyb[0:3])+sum(bc[0:3])
                    # io.imsave(folderlist[i]+'cyto.tif',im1)
        
        cyto_sum_gaussian=filters.gaussian(cyto_sum,sigma=5,preserve_range=True)
        
                    # io.imsave(folderlist[i]+'cyto_gaussian.tif',im1)
        
        
                    # im2=im0[4,:,:]
        nuclei=hyb[4]
        # im2=filters.gaussian(im2,sigma=3,preserve_range=False)
        #    plt.figure()
        #    plt.imshow(im1)
        combined_img=np.concatenate((cyto_sum_gaussian[:,:,None],nuclei[:,:,None]),axis=2)
        # im3=np.concatenate((im1[:,:,None],im2[:,:,None]),axis=2)
        imgs.append(combined_img)
        # define CHANNELS to run segementation on
        # grayscale=0, R=1, G=2, B=3
        # channels = [cytoplasm, nucleus]
        # if NUCLEUS channel does not exist, set the second channel to 0
        channels = [1,2]
                    
                    
        masks, flows, styles, diams = model.eval(imgs, diameter=50, channels=channels,
            flow_threshold=0.6, cellprob_threshold=0)
        
        maski = masks[0]
        scipy.io.savemat(folderlist[i]+'cellmask.mat',dict(maski=maski))


# In[3]: