@author: Keegstra

What this function does:

1. load images from folder, as well as auxiliary files and images.
2. do background and flatfield correction
3. align them by means of image cross-correlation
4. find seeds of cells
5. filter seeds on size and distance to border.
6. construct masks and filter masks for overlap.
6. find intensity per cell
7. do bleaching correction
8. save all good cells in a file.

"""

#====================++++++++++++++++++++++++++++++++++++++++==================
#
#                                 IMPORT DEPENDENCIES
#
#===================+++++++++++++++++++++++++++++++++++++++++==================

import numpy as np
import math
import scipy.signal as signal
import scipy.optimize as opt
import scipy.io as io
import mahotas as mh
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import os
import sys
import csv
import time
plt.ion()

#===========================+++++++++++++++++++++++++++========================
#
#                                DEFINE PARAMETERS
#
#+++++++++++++++++++++++++++===========================++++++++++++++++++++++++
myinput = sys.argv[1]
print("Your input was: " + myinput)
#folderbase="D:/Single_cell_fret/JACKSON_DATA_2/Utah/TSR_F396L/SERIES2_20150618"
folderbase=myinput
maxSize=75 #Standard 75, otherwise risk of selectinf doubles.
minSize=6
maxMargin=8;
maxNoverlap=4
w=17
maxOverlap=4*w
threshold=200 #threshold for max RMS error
red_gating=1000;
th=800
BCKval=900;


#====================++++++++++++++++++++++++++++++++++++++++==================
#
#                    DEFINE FUNCTIONS
#
#===================+++++++++++++++++++++++++++++++++++++++++==================

w=w-1;

def makePaths(folder):
    fileFolder=folderbase+"/Registered/Frames000001/"
    fileFolderRED=folderbase+"/RED/Frames000001"
    BCKpath=folderbase+"/BCK.tif"
    stimuluspath=folderbase+"/stimuli.csv"
    savepath=folderbase+"/segmentation_data.mat"
    FF_REDpath=folderbase+"/FF_RED.tif"
    FF_GREENpath=folderbase+"/FF_GREEN.tif"
    return fileFolder,fileFolderRED,BCKpath,stimuluspath,savepath,FF_REDpath,FF_GREENpath

def loadImageStack(folder):
    filenames_tif=[file for file in os.listdir(folder) if file.endswith(".Tiff")]
    if len(filenames_tif)==0:
        filenames_tif=[file for file in os.listdir(folder) if file.endswith(".tif")] 
        
    nimg=len(filenames_tif)
    image_stack=np.zeros((512,512,nimg))
    for i in range(nimg):
        image_stack[:,:,i]=mh.imread(folder+'/'+filenames_tif[i])
        #then write progress in command line.
        sys.stdout.write("\r loading files :"+str((i+1))+"/"+str((nimg)))
        sys.stdout.flush()
    image_stack=image_stack.astype(np.uint16)
    return image_stack,nimg
    
def read_stimuli(path):
    rows=[]; index_neg=[]; index_pos=[]
    try:
        with open(path) as comments:
            reader=csv.reader(comments)
            for row in reader:
                rows.append(row)
                index_neg.append(row[1].find('add',0,3))
                index_pos.append(row[1].find('rem',0,3))
            neg_inds_pos=[i for i,item in enumerate(index_neg) if item>-1]
            neg_inds=[]; 
        for n in neg_inds_pos:
            neg_inds.append(int(float(rows[n][0])))
        pos_inds=[]
        pos_inds_pos=[i for i,item in enumerate(index_pos) if item>-1]
        for n in pos_inds_pos:
            pos_inds.append(int(float(rows[n][0])))        
    except IOError:
       neg_inds=[]; pos_inds=[] 
       print '\n Unable to load stimuli data.'
  
    return neg_inds,pos_inds

def correctBackground(image_stack_raw,BCKpath,BCKval):
    nimg=len(image_stack_raw[1,1,:])
    image_stack=np.zeros(image_stack_raw.shape).astype(np.uint16)    
    try:
        BCK=mh.imread(BCKpath)
        for i in range(nimg):
            corrected=image_stack_raw[:,:,i].astype('double')-BCK
            image_stack[:,:,i]=corrected.clip(0).astype(np.uint16)
        report=1;
    except IOError:
        print('\n Unable to load BCK image. Taking uniform standard value of {} instead.'.format(BCKval))
        BCK=BCKval
        image_stack[image_stack_raw>BCK]=image_stack_raw[image_stack_raw>BCK]-BCK
        report=0;
    return image_stack,report


def loadAcceptorChannel(folderRED,BCKpath,image_stack):
    try:
        image_stack_rawRED,nimgRED=loadImageStack(folderRED)
        image_stackRED,report=correctBackground(image_stack_rawRED,BCKpath,BCKval)
        ACCim=np.mean(image_stackRED,axis=2)
        yoffR,xoffR=align_channels(image_stack[:,:,1],ACCim)
        if (yoffR>=0)&(xoffR<=0):
            ACC=ACCim[0:(511-yoffR),(256-xoffR):511]
        elif (yoffR<=0)&(xoffR<=0):
            ACC=ACCim[(0-yoffR):511,(256-xoffR):511]
        elif (yoffR>=0)&(xoffR>=0):
            ACC=ACCim[0:(511-yoffR),(256+xoffR):511]
        elif (yoffR<=0)&(xoffR>=0):
            ACC=ACCim[(0-yoffR):511,(256+xoffR):511]
    except WindowsError:
        print '\n Unable to load acceptor channel images.'   
        ACCim=np.zeros(image_stack[:,:,1].shape)
        xoffR=0
        yoffR=0
        ACC=ACCim[0:(511-yoffR),(256-xoffR):511]
    return ACC,xoffR,yoffR

def align_channels(image1,image2):
    sample1=image1[200:300,100:200]-np.mean(image1[200:300,100:200])
    sample2=image2[200:300,356:456]-np.mean(image2[200:300,356:456])
    corr=signal.correlate2d(sample1,sample2); #it makes it faster to do the calculation for a small area.
    i,j=np.unravel_index(corr.argmax(),corr.shape)
    yoff=i-99;
    xoff=j-99;
    return yoff,xoff
      
def split_channels(image_stack):
    yoff,xoff=align_channels(image_stack[:,:,1],image_stack[:,:,1]);
    #greenstack=image_stack[0:511,0:255,:]
    #redstack=image_stack[0:(511),(256):511,:]
    if (yoff>=0)&(xoff<=0):
        greenstack=image_stack[(0+yoff):511,0:(255+xoff),:]
        redstack=image_stack[0:(511-yoff),(256-xoff):511,:]
    elif (yoff<=0)&(xoff<=0):
        greenstack=image_stack[0:(511+yoff),0:(255+xoff),:]
        redstack=image_stack[(0-yoff):511,(256-xoff):511,:]
    elif (yoff>=0)&(xoff>=0):
        greenstack=image_stack[(0+yoff):511,(0+xoff):255,:]
        redstack=image_stack[0:(511-yoff),256:(511-xoff),:]
    elif (yoff<=0)&(xoff>=0):
        greenstack=image_stack[0:(511+yoff),(0+xoff):255,:]
        redstack=image_stack[(0-yoff):511,256:(511-xoff),:]
        #redstack=image_stack[0:(500-yoff),(256+xoff):511,:]
    return greenstack,redstack,xoff,yoff

def correctFlatfield(greenstack,redstack,ACC,folderbase,xoff,yoff,xoffR,yoffR):
    FF_REDpath=folderbase+"/FF_RED.tif"
    FF_GREENpath=folderbase+"/FF_GREEN.tif"   
    try:    
        FF_RED=mh.imread(FF_REDpath)
        FF_GREEN=mh.imread(FF_GREENpath)
        report=1;
    except IOError:
        print '\n Unable to do flatfield correction'
        FF_RED=np.ones((512,512))
        FF_GREEN=np.ones((512,512))
        report=0;
    if (yoff>=0)&(xoffR<=0):
        FFg=FF_GREEN[(0+yoff):511,0:(255+xoff)]
        FFr=FF_RED[0:(511-yoffR),(256-xoffR):511]
    elif (yoff<=0)&(xoffR<=0):
        FFg=FF_GREEN[0:(511+yoff),0:(255+xoff)]
        FFr=FF_RED[(0-yoffR):511,(256-xoffR):511]
    elif (yoff>=0)&(xoffR>=0):
        FFg=FF_GREEN[(0+yoff):511,(0+xoff):255]
        FFr=FF_RED[0:(511-yoffR),256:(511-xoffR)]        
    elif (yoff<=0)&(xoffR>=0):
        FFg=FF_GREEN[0:(511+yoff),(0+xoff):255]
        FFr=FF_RED[(0-yoffR):511,256:(511-xoffR)]
    #normalize flat-field images
    FFg=FFg.astype('double')/np.max(FFg)
    FFr=FFr.astype('double')/np.max(FFr)
    nimg=len(greenstack[1,1,:])
    for i in range(nimg):
        corrected=np.divide(greenstack[:,:,i].astype('double'),FFg)
        greenstack[:,:,i]=corrected.clip(0).astype(np.uint16)
        corrected=np.divide(redstack[:,:,i].astype('double'),FFg)
        redstack[:,:,i]=corrected.clip(0).astype(np.uint16)    
    corrected=np.divide(ACC[:,:].astype('double'),FFr)
    ACC[:,:]=corrected.clip(0).astype(np.uint16)    
    return greenstack,redstack,ACC,report   


def get_labeled_cell_coordinates(labeled,ncells):
    bboxes = mh.labeled.bbox(labeled)
    xcor=[0]*ncells
    ycor=[0]*ncells
    seeds_img=np.zeros(shape=labeled.shape)
    for i in [item +1 for item in range(ncells)]:
        inds=[item for item in bboxes[i]]
        ycor[i-1]=np.round(np.mean(inds[0:2]))
        xcor[i-1]=np.round(np.mean(inds[2:4]))
        seeds_img[ycor[i-1],xcor[i-1]]=1;
    return xcor,ycor,seeds_img
    
def mask_overlap(x_start,x_end,y_start,y_end,mask_img):
    border=[]
    borders=[];
    noverlaps=[0]*ncells
    x_range=[0]*ncells
    y_range=[0]*ncells
    for i in range(ncells): 
        y_range[i]=np.linspace(y_start[i],y_end[i]-1,w).astype(int)
        x_range[i]=np.linspace(x_start[i],x_end[i]-1,w).astype(int)
        border.extend(mask_img[y_range[i],x_range[i][0]]) 
        border.extend(mask_img[y_range[i][-1],x_range[i]]) 
        border.extend(mask_img[np.fliplr([y_range[i][1:]])[0],x_range[i][-1]]) 
        border.extend(mask_img[y_range[i][0],np.fliplr([x_range[i][1:]])[0]])
        #border.append(total_matrix[y_range[i][0],x_range[i][0]]) #append the fist value because of circular boundary condition
        borders.append(border)
        border=[]
        noverlaps[i]=np.sum(np.diff(borders[i]).clip(0))     
    sumoverlap=[0]*ncells
    for n in range(ncells): 
            sumoverlap[n]=np.sum(mask_img[y_range[n][0]:y_range[n][-1],x_range[n][0]:x_range[n][-1]])-(w)**2

    return noverlaps,sumoverlap,borders

def build_masks(labeled,ncells,xcor,ycor):
    a,b=labeled.shape
    masks=np.zeros(shape=(a,b,ncells))
    masksLabeled=np.zeros(shape=(a,b,ncells))
    x_start=[0]*ncells;
    x_end=[0]*ncells;
    y_start=[0]*ncells;
    y_end=[0]*ncells;
    for i in range(ncells):
        x_start[i]=xcor[i]-math.ceil(w/2) 
        x_end[i]=xcor[i]+math.ceil(w/2)+1
        y_start[i]=ycor[i]-math.ceil(w/2)
        y_end[i]=ycor[i]+math.ceil(w/2)+1
        masks[y_start[i]:y_end[i],x_start[i]:x_end[i],i]=1
        masksLabeled[y_start[i]:y_end[i],x_start[i]:x_end[i],i]=(i+1)
    mask_img=masks.sum(2)
    maskLabeled_img=masksLabeled.sum(2)
    return masks,masksLabeled,mask_img,maskLabeled_img,x_start,x_end,y_start,y_end

def build_segmentation_boxes(labeled,ncells,w):
    xcor,ycor,seeds_img=get_labeled_cell_coordinates(labeled,ncells)
    masks,masksLabeled,mask_img,maskLabeled_img,x_start,x_end,y_start,y_end=build_masks(labeled,ncells,xcor,ycor)
    noverlaps,sums,borders=mask_overlap(x_start,x_end,y_start,y_end,mask_img)
    return x_start,x_end,y_start,y_end,xcor,ycor,borders,seeds_img,masks,mask_img,maskLabeled_img,noverlaps,sums,

def filter_cell_size(labeled,minSize,maxSize):
    sizes = mh.labeled.labeled_size(labeled)
    too_big = np.where(sizes > maxSize)
    labeled = mh.labeled.remove_regions(labeled, too_big)
    too_small = np.where(sizes < minSize)
    relabeled = mh.labeled.remove_regions(labeled, too_small)  
    relabeled, ncells = mh.labeled.relabel(relabeled)   
    return relabeled,ncells
 
def filter_cell_proximity(labeled,noverlaps,sums,maxNoverlap,maxOverlap):
    too_close = np.where(np.asarray(noverlaps) > maxNoverlap)
    too_close = [item+1 for item in too_close]
    labeled = mh.labeled.remove_regions(labeled, too_close)
    #relabeled, ncells = mh.labeled.relabel(labeled)   
    still_too_close = np.where(np.asarray(sums) > maxOverlap)
    still_too_close = [item+1 for item in still_too_close]
    labeled = mh.labeled.remove_regions(labeled, still_too_close)
    relabeled, ncells = mh.labeled.relabel(labeled)   
    return relabeled,ncells

def filter_acceptor_intensities(Iacc,red_gating,good_inds):
    bad_inds=[]
    if sum(Iacc)>0:
        for n in good_inds:
            if Iacc[n]<red_gating:
                bad_inds.append(n)
    s=set(bad_inds)
    good_inds=[item for item in good_inds if item not in s]
    print('\n {} cells were excluded because acceptor intensity was too low.'.format(len(bad_inds)))
    return good_inds,bad_inds
   
def remove_border_cells(labeled,maxMargin):
    labeled = mh.labeled.remove_bordering(labeled,maxMargin)
    relabeled, ncells = mh.labeled.relabel(labeled)   
    return relabeled,ncells

def add_rectangle_to_plot(axes,x_pos,y_pos,width):
    for n in range(len(x_pos)):     
        axes.add_patch(
            patches.Rectangle(
            (x_pos[n]-math.ceil(w/2), y_pos[n]-math.ceil(w/2)),width,width,
            fill=False, edgecolor="red", linewidth=3
            )
        )

def get_fluorescence_intensities(greenstack,redstack,x_start,x_end,y_start,y_end,ncells,nimg):
    Igreen=np.zeros((ncells,nimg))
    Ired=np.zeros((ncells,nimg))
    R=np.zeros((ncells,nimg))
    for i in range(nimg):
        img_green=greenstack[:,:,i]
        img_red=redstack[:,:,i]
        for n in range(ncells):
            Igreen[n,i]=np.sum(img_green[y_start[n]:y_end[n],x_start[n]:x_end[n]])
            Ired[n,i]=np.sum(img_red[y_start[n]:y_end[n],x_start[n]:x_end[n]])
            #Igreen[n,i]=np.sum(img_green[(ycor[n]-math.ceil(w/2)-a):(ycor[n]+math.ceil(w/2)-a),(xcor[n]-math.ceil(w/2)-b):(xcor[n]+math.ceil(w/2)-b)])
            #Ired[n,i]=np.sum(img_red[(ycor[n]-math.ceil(w/2)-a):(ycor[n]+math.ceil(w/2)-a),(xcor[n]-math.ceil(w/2)-b):(xcor[n]+math.ceil(w/2)-b)])
        R[:,i]=Ired[:,i]/Igreen[:,i]
    return Ired,Igreen,R

def get_fluorescence_intensity(image_stack,x_start,x_end,y_start,y_end,ncells,nimg):
    Icell=np.zeros((ncells,nimg))
    if nimg<5:
        IcellIm=np.zeros((ncells,w+1,w+1))
        for n in range(ncells):
            IcellIm[n,:,:]=image_stack[y_start[n]:y_end[n],x_start[n]:x_end[n]]
    else:
        IcellIm=np.zeros((w+1,w+1,nimg,ncells))
        for n in range(ncells):
            IcellIm[:,:,:,n]=image_stack[y_start[n]:y_end[n],x_start[n]:x_end[n],:]
 
    for i in range(nimg):
        if nimg > 1:
            img=image_stack[:,:,i]
        elif  nimg == 1:
            img=image_stack
        for n in range(ncells):
            Icell[n,i]=np.sum(img[y_start[n]:y_end[n],x_start[n]:x_end[n]])
    return Icell,IcellIm

def build_fit_segment(neg_inds,pos_inds,nimg):
    if len(neg_inds)>2:
        fitsegment=[]
        for neg in neg_inds:
            fitsegment=np.concatenate((fitsegment,(np.arange((neg-10),(neg+5))))).astype(int)
        fitsegment=np.concatenate((fitsegment,np.arange((nimg-20),nimg))).astype(int)
    else: 
        fitsegment=np.arange(1,nimg).astype(int)
    return fitsegment    
        
    
def single_exponential(x,a,b,c):
    return a*np.exp(b * x)+c

def simple_linear(x,a,b):
    return a*x+b

def rmserror(t,data,fitfunc,pars):
    rms=np.square(np.sum(np.square((data-fitfunc(t,*pars))/np.mean(data))))
    return rms
  
def fit_bleaching_decay(t,fitsegment,Ired,Igreen,threshold,fitfunc):
    ncells=len(Igreen[:,1])
    Igcor=np.zeros(Igreen.shape)
    Ircor=np.zeros(Igreen.shape)
    R=np.zeros(Igreen.shape)
    popt=[1]*3
    predopt=[1]*3
    inds=fitsegment
    fitgreen=[]
    fitred=[]
    bad_inds=[]
    for n in range(ncells):
        if fitfunc == single_exponential:
            pguessgreen=[ Igreen[n,np.arange(25)].mean()-Igreen[n,-25:].mean(),-0.001, Igreen[n,-25:].mean()]
            pguessred=[ Ired[n,np.arange(25)].mean()-Ired[n,-25:].mean(),-0.0005, Ired[n,-25:].mean()]
        elif fitfunc == simple_linear: 
            pguessgreen=[ Igreen[n,np.arange(25)].mean()-Igreen[n,-25:].mean(), Igreen[n,-25:].mean()]
            pguessred=[ Ired[n,np.arange(25)].mean()-Ired[n,-25:].mean(), Ired[n,-25:].mean()] 
        try:
            popt,pcov=opt.curve_fit(fitfunc,t[inds],Igreen[n,inds],p0=pguessgreen)
            predopt,predcov=opt.curve_fit(fitfunc,t[inds],Ired[n,inds],p0=pguessred)
            fitgreen.append(popt) 
            fitred.append(predopt)
            Igcor[n,:]=(fitfunc(0,*popt)) * Igreen[n,:] / fitfunc(t,*popt)
            Ircor[n,:]=(fitfunc(0,*predopt)) * Ired[n,:] / fitfunc(t,*predopt)
        except RuntimeError:
            Igcor[n,:]=Igreen[n,:]
            Ircor[n,:]=Ired[n,:]        
            fitgreen.append(popt) 
            fitred.append(predopt)
            bad_inds.append(n)
        R[n,:]=Ircor[n,:]/Igcor[n,:]  
    
    s=set(bad_inds)
    good_inds=[item for item in range(ncells) if item not in s]

    rms_g=[0]*ncells
    for n in good_inds:
        rms_g[n]=rmserror(t,Igreen[n,:],fitfunc,fitgreen[n])
    wrong_fits=[item for item, ind in enumerate(rms_g) if ind > threshold]
    print('\n {} cells were excluded because unable to fit.'.format(len(bad_inds)))

    s=set(wrong_fits)
    good_inds=[item for item in good_inds if item not in s]
    print('\n {} cells were excluded because fit was bad.'.format(len(wrong_fits)))
    print('\n Remaining are {} cells.'.format(len(good_inds)))
    print '\n Bleaching correction completed.'
    return Ircor,Igcor,R,fitred,fitgreen,good_inds,bad_inds,wrong_fits
    
          
#==============================================================================
#
#                                  CODE TO EXECUTE
#      
#==============================================================================

plt.close('all')      

#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+    
#PRELIMINARIES: LOAD DATA, AUXILLIARY FILES,
#DO BASIC CORRECTIONS 
#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+   

folder,folderRED,BCKpath,stimuluspath,savepath,FF_REDpath,FF_GREENpath=makePaths(folderbase)

#read in the files
image_stack_raw,nimg=loadImageStack(folder)

##background correction
image_stack,report_background=correctBackground(image_stack_raw,BCKpath,BCKval)

#read in comments
neg_inds,pos_inds=read_stimuli(stimuluspath)

##split red and green channels
greenstack,redstack,xoff,yoff=split_channels(image_stack)

#load red channel data
ACC,xoffR,yoffR=loadAcceptorChannel(folderRED,BCKpath,image_stack)


#do flat-field correction
greenstack,redstack,ACC,report_flatfield=correctFlatfield(greenstack,redstack,ACC,folderbase,xoff,yoff,xoffR,yoffR)

#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+    
#SEGMENTATION
#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+   

fig3=plt.figure( figsize = (15,7) )
ax1=fig3.add_subplot(151)
plt.imshow(greenstack[:,:,0], 'gray')
plt.title('Raw')

fig3.add_subplot(152)
th = mh.thresholding.otsu(greenstack[:,:,0])
greenB_th = greenstack[:,:,1] > ( th )
plt.imshow(greenB_th, 'gray')
plt.title('Threshold')

fig3.add_subplot(153)
labeled, nlbl = mh.label(greenB_th)
labeled,ncells=filter_cell_size(labeled,minSize,maxSize)
labeled,ncells=remove_border_cells(labeled,maxMargin)
plt.imshow(labeled)
plt.title('Segmented')

ax4=fig3.add_subplot(154)
x_start,x_end,y_start,y_end,xcor,ycor,borders,seeds_img,masks,mask_img,maskLabeled_img,noverlaps,sums=build_segmentation_boxes(labeled,ncells,w)
plt.imshow(mask_img,'gray')
add_rectangle_to_plot(ax4,xcor,ycor,w)
plt.title('All ROIs')

ax5=plt.subplot(155)
labeled,ncells=filter_cell_proximity(labeled,noverlaps,sums,maxNoverlap,maxOverlap)
x_start,x_end,y_start,y_end,xcor,ycor,borders,seeds_img,masks,mask_img,maskLabeled_img,noverlaps,sums=build_segmentation_boxes(labeled,ncells,w)
plt.imshow(mask_img, 'gray')
plt.title('ROIs after filtering')


ax1=plt.subplot(151)
add_rectangle_to_plot(ax1,xcor,ycor,w)
plt.show()

#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+    
#GET THE SIGNAL FROM THE IMAGE STACK
#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+   

Igreen,IgreenIm=get_fluorescence_intensity(greenstack,x_start,x_end,y_start,y_end,ncells,nimg)
_,SegIm=get_fluorescence_intensity(greenB_th,x_start,x_end,y_start,y_end,ncells,1)
Ired,IredIm=get_fluorescence_intensity(redstack,x_start,x_end,y_start,y_end,ncells,nimg)
Iacc,IaccIm=get_fluorescence_intensity(ACC,x_start,x_end,y_start,y_end,ncells,1)

#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+    
# DO BLEACHING CORRECTION
#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+   

t=np.arange(nimg)
tfit=build_fit_segment(neg_inds,pos_inds,nimg)
fitfunc=simple_linear#single_exponential#simple_linear
Ircor,Igcor,R,fitred,fitgreen,good_inds,bad_inds,wrong_fits=fit_bleaching_decay(t,tfit,Ired,Igreen,threshold,fitfunc)

#filter cells with too low CheY expression.
good_inds,low_acceptor=filter_acceptor_intensities(Iacc,red_gating,good_inds)

#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+    
# FILTER & ORGANIZE DATA
#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+   

Igreen_arranged=np.zeros((len(good_inds),Igreen.shape[1]))
IgreenIm_arranged=np.zeros((w+1,w+1,Igreen.shape[1],len(good_inds)))
SegIm_arranged=np.zeros((w+1,w+1,len(good_inds)))
Ired_arranged=np.zeros((len(good_inds),Ired.shape[1]))
IredIm_arranged=np.zeros((w+1,w+1,len(good_inds)))
Iacc_arranged=np.zeros((len(good_inds),Iacc.shape[1]))
IaccIm_arranged=np.zeros((w+1,w+1,len(good_inds)))
Igcor_arranged=np.zeros((len(good_inds),Igcor.shape[1]))
Ircor_arranged=np.zeros((len(good_inds),Ircor.shape[1]))
R_arranged=np.zeros((len(good_inds),R.shape[1]))
xcor_arranged=np.zeros((len(good_inds),1))
ycor_arranged=np.zeros((len(good_inds),1))
for n in range(len(good_inds)):
    Igcor_arranged[n,:]=Igcor[good_inds[n],:]
    Ircor_arranged[n,:]=Ircor[good_inds[n],:]
    Igreen_arranged[n,:]=Igreen[good_inds[n],:]
    Ired_arranged[n,:]=Ired[good_inds[n],:]
    IgreenIm_arranged[:,:,:,n]=IgreenIm[:,:,:,good_inds[n]]
    IaccIm_arranged[:,:,n]=IaccIm[good_inds[n],:,:]
    SegIm_arranged[:,:,n]=SegIm[good_inds[n],:,:]
    Iacc_arranged[n,:]=Iacc[good_inds[n],:]
    R_arranged[n,:]=R[good_inds[n],:]
    xcor_arranged[n,]=xcor[good_inds[n]]
    ycor_arranged[n]=ycor[good_inds[n]]
fitgreen_arranged=[fitgreen[n] for n in good_inds]
fitred_arranged=[fitred[n] for n in good_inds]    
         
          


#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+    
# SAVE DATA
#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 

metadata = {'w':w,'maxOverlap':maxOverlap,'maxNoverlap':maxNoverlap,'maxSize':maxSize,'minSize':minSize,'maxMargin':maxMargin}
analysis_report = {'ncells':len(good_inds),'excluded_wrong_fit': len(wrong_fits),'excluded_no_fit':len(bad_inds),'excluded_low_acceptor':len(low_acceptor),'background_correction':report_background,'flatfield_correction':report_flatfield}
#data = {'Ired': Ired_arranged,'Igreen': Igreen_arranged,'Igcor': Igcor_arranged,'Ircor': Ircor_arranged,'IgreenIm':IgreenIm_arranged,'IaccIm':IaccIm_arranged,'SegIm':SegIm_arranged,'xcor':xcor_arranged,'ycor':ycor_arranged,'image':greenstack[:,:,1],'pos_inds':pos_inds,'neg_inds':neg_inds,'metadata':metadata,'analysis_report':analysis_report}
data = {'Ired': Ired_arranged,'Igreen': Igreen_arranged,'Igcor': Igcor_arranged,'Ircor': Ircor_arranged,'xcor':xcor_arranged,'ycor':ycor_arranged,'image':greenstack[:,:,1],'pos_inds':pos_inds,'neg_inds':neg_inds,'metadata':metadata,'analysis_report':analysis_report}

w=w+1;
io.savemat(savepath,data)


#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+    
# PLOT THINGS
#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 

fig2=plt.figure( figsize = (15,7) )
green=greenstack[:,:,1]
red=redstack[:,:,1]
ax1 = fig2.add_subplot(131)
plt.imshow(green, 'gray')
add_rectangle_to_plot(ax1,xcor,ycor,w)
plt.title('Donor FRET intensity')

ax2 = fig2.add_subplot(132)
plt.imshow(red, 'gray')
add_rectangle_to_plot(ax2,xcor,ycor,w)
plt.title('Acceptor FRET intensity')

ax3 = fig2.add_subplot(133)
plt.imshow(ACC, 'gray')
add_rectangle_to_plot(ax3,xcor,ycor,w)
plt.title('Acceptor direct excitation')

plt.suptitle('Red and Green Channels')

plt.show()



tplot = np.empty(nimg)
tplot.fill(np.nan)
tplot[tfit]=tfit

fig3=plt.figure( figsize = (20,10) )
fig3.subplots_adjust(wspace=0.25,hspace=0.6)
plt.clf()
ax1 = plt.subplot2grid((6,3), (2,0), colspan=2,rowspan=2)
green_plt, =ax1.plot(t,Igreen[1,:],'k-')
green_corr_plt, =ax1.plot(t,Igcor[1,:],'g-')
green_used_plt, =ax1.plot(tplot,Igreen[1,:],'b-')
green_fit_plt, =ax1.plot(t,fitfunc(t,*fitgreen[1]),'b-')
plt.ylabel('D')

ax2 = plt.subplot2grid((6,3), (0,0), colspan=2,rowspan=2)
red_plt, =ax2.plot(t,Ired[1,:],'k-')
red_corr_plt, =ax2.plot(t,Ircor[1,:],'r-')
red_used_plt, =ax2.plot(tplot,Ired[1,:],'b-')
red_fit_plt, =ax2.plot(t,fitfunc(t,*fitred[1]),'b-')
plt.ylabel('A')

ax3 = plt.subplot2grid((6,3), (4,0), colspan=2,rowspan=2)
ratio_plt,=ax3.plot(t,Ircor[1,:]/Igcor[1,:],'k-')
plt.ylabel('A/D')
plt.xlabel('time')
 
ax5 = plt.subplot2grid((6,3), (3,2),rowspan=3)
Iacc_plt,=ax5.plot(Igreen[1,1],Iacc[1],'ro')
Ired_plt,=ax5.plot(Igreen[1,1],Ired[1,1],'ko')
plt.xlabel('Green Intensity'); plt.ylabel('Red intensity')
plt.xlim((0,Igreen[:,1].max())); plt.ylim((0,Ired[:,1].max()))

for i in good_inds:
    green_plt.set_ydata(Igreen[i,:])
    green_corr_plt.set_ydata(Igcor[i,:])
    green_used_plt.set_ydata(Igreen[i,:])
    green_fit_plt.set_ydata(fitfunc(t,*fitgreen[i]))
    ax1.relim()
    ax1.autoscale_view(True,True,True)

    red_plt.set_ydata(Ired[i,:])
    red_corr_plt.set_ydata(Ircor[i,:])
    red_used_plt.set_ydata(Ired[i,:])
    red_fit_plt.set_ydata(fitfunc(t,*fitred[i]))
    ax2.relim()
    ax2.autoscale_view(False,True,True)

    ratio_plt.set_ydata(Ircor[i,:]/Igcor[i,:])
    ax3.relim()
    ax3.autoscale_view(True,True,True)

    Iacc_plt.set_xdata(Igreen[i,1])
    Iacc_plt.set_ydata(Iacc[i])
    Ired_plt.set_xdata(Igreen[i,1])
    Ired_plt.set_ydata(Ired[i,1])

 
    plt.draw()

plt.ioff()
ax4 = plt.subplot2grid((6,3), (0,2),rowspan=3)
plt.plot(t,np.mean(Ircor_arranged/Igcor_arranged,axis=0),color='r',linewidth=1.0)
plt.plot(t,np.mean(Ired/Igreen,axis=0),color='k')
plt.legend(['corrected','uncorrected'],loc='best')
plt.title('population averaged')

ax5 = plt.subplot2grid((6,3), (3,2),rowspan=3)

plt.plot(Igreen_arranged[:,1],Iacc_arranged,'ro')
plt.plot(Igreen_arranged[:,1],Ired_arranged[:,1],'ko')
plt.xlabel('Green Intensity'); plt.ylabel('Red intensity')
plt.legend(['Independent','from FRET'],loc='best')
plt.title('donor vs acceptor')

ax1 = plt.subplot2grid((6,3), (0,0), colspan=2,rowspan=6)

for row in R_arranged: 
    plt.plot(t,row)
plt.plot(t,np.mean(R_arranged,axis=0),color='k',linewidth=3.0)
plt.xlabel('time (frame)'); plt.ylabel('acceptor/donor'); plt.title('single cell time traces')
plt.show()
#plt.subplot(2,3,3)

print("\n goodbye")



