import colonyzer2 as c2
import json
import argparse
import shutil
import string
import os
import time
import numpy as np
import itertools
import cv2
from matplotlib.backends.backend_pdf import PdfPages, FigureCanvasPdf
from matplotlib import figure
from PIL import Image,ImageDraw
from scipy import ndimage
from pkg_resources import resource_filename, Requirement

def checkImages(fdir,fdict=None,barcRange=(0,-24),verbose=False):
    '''Discover barcodes in current working directory (or in fdir or in those specified in fdict) for which analysis has not started.'''
    if fdict!=None:
        with open(fdict, 'rb') as fp:
            barcdict = json.load(fp)
            # Drop any barcodes that are currently being analysed/already analysed
            barcdict={x:barcdict[x] for x in barcdict.keys() if not c2.checkAnalysisStarted(barcdict[x][-1])}
    else:
        # Find image files which have yet to be analysed
        # Lydall lab file naming convention (barcRange)
        # First 15 characters in filename identify unique plates
        # Remaining charaters can be used to store date, time etc.
        barcdict=c2.getBarcodes(fdir,barcRange,verbose=verbose)
    return(barcdict)

def parseArgs(inp=''):
    '''Define console script behaviour, hints and documentation for setting off Colonyzer analysis.'''
    parser=argparse.ArgumentParser(description="Analyse timeseries of QFA images: locate cultures on plate, segment image into agar and cells, apply lighting correction, write report including cell density estimates for each location in each image.  If you need to specify initial guesses for colony locations, you must provide a Colonyzer.txt file (as generated by ColonyzerParametryzer) describing initial guess for culture array in the directory containing the images to be analysed.")

    parser.add_argument("-c","--lc", help="Enable lighting correction?", action="store_true")
    parser.add_argument("-m","--diffims", help="If lighting correction switched on, attempt to correct for lighting differences between images in timecourse (can induce slight negative cell density estimates).", action="store_true")
    parser.add_argument("-p","--plots", help="Plot pixel intensity distributions, segmentation thresholds and spot location traces?", action="store_true")
    parser.add_argument("-i","--initpos", help="Use intial guess for culture positions from Colonyzer.txt file?", action="store_true")
    parser.add_argument("-x","--cut", help="Cut culture signal from first image to make pseudo-empty plate?", action="store_true")
    parser.add_argument("-q","--quiet", help="Suppress messages printed to screen during analysis?", action="store_true")
    parser.add_argument("-e","--endpoint", help="Only analyse final image in series.  Mostly for testing single image analysis.",action="store_true")
    parser.add_argument("-k","--edgemask", help="Use intensity gradient & morphology for image segmentation instead of thresholding.",action="store_true")
    parser.add_argument("-g","--greenlab", help="Check for presence of GreenLab lids on plates.",action="store_true")
    
    parser.add_argument("-d","--dir", type=str, help="Directory in which to search for image files that have not been analysed (current directory by default).",default=".")
    parser.add_argument("-l","--logsdir", type=str, help="Directory in which to search for JSON files listing images for analyis (e.g. LOGS3, root of HTS filestore).  Only used when intending to specify images for analysis in .json file (see -u).",default=".")
    parser.add_argument("-f","--fixthresh", type=float, help="Image segmentation threshold value (default is automatic thresholding).")
    parser.add_argument("-u","--usedict", type=str, help="Load .json file specifying images to analyse.  If argument has a .json extension, treat as filename.  Otherwise assume argument is a HTS-style screen ID and return path to appropriate .json file from directory structure.  See C2Find.py in HTSauto package.")
    parser.add_argument("-o","--fmt", type=str, nargs='+', help="Specify rectangular grid format, either using integer shorthand (e.g. -o 96, -o 384, -o 768 -o 1536) or explicitly specify number of rows followed by number of columns (e.g.: -o 24 16 or -o 24x16).", default=['384'])
    parser.add_argument("-t","--updates", type=int, help="Number of (quasi-)randomly distributed grid positions to assess in first phase of grid location.  Ignored when -initpos specified.", default=144)
    
    if inp=="":
        args = parser.parse_args()
    else:
        args = parser.parse_args(inp.split())
    return(args)

def buildVars(inp=''):
    '''Read user input, set up flags for analysis, report on options chosen and find files to be analysed.'''
    inp=parseArgs(inp)

    if inp.quiet:
        verbose=False
    else:
        verbose=True
    
    if inp.dir==None:
        fdir=os.getcwd()
    else:
        fdir=os.path.realpath(inp.dir)

    if inp.fixthresh!=None and not inp.edgemask:
        fixedThresh=inp.fixthresh
    else:
        fixedThresh=-99

    if len(inp.fmt)>2:
        print("Woah!  Too many dimensions specified for rectangular grid format!")
        nrow,ncol=(0,0)
    elif len(inp.fmt)==1:
        nrow,ncol=c2.parsePlateFormat(inp.fmt[0])
    else:
        nrow,ncol=[int(x) for x in inp.fmt]
    print("Expecting {} rows and {} columns on plate.".format(nrow,ncol))
    
    if inp.usedict is None:
        fdict=None    
    elif inp.usedict[-5:] in [".json",".JSON"]:
        fdict=os.path.realpath(inp.usedict)
    else:
        fdict=locateJSON(inp.usedict,os.path.realpath(inp.logsdir),verbose)
    if fdict is not None and not os.path.exists(fdict): print("WARNING! "+fdict+" does not exist...")

    if inp.lc:
        diffIms=True
    else:
        diffIms=False

    if inp.lc and inp.cut:
        cut=True
    else:
        cut=False

    if verbose:
        if inp.lc:
            print("Lighting correction turned on.")
        else:
            print("Lighting correction turned off.")
        if inp.lc:
            if inp.diffims:
                print("Correcting for lighting differences between subsequent images of same plate.")
            else:
                print("Any lighting differences between plates will be ignored.")
        if inp.plots:
            print("Reports on spot location and thresholding will appear in Output_Reports directory.")
        else:
            print("No reports on spot location or thresholding will be generated.")
        if inp.initpos:
            print("Using user-specified initial guess for colony locations.  NOTE: Colonyzer.txt file must be located in directory with images to be analysed.  See Parametryzer for more information.")
        else:
            print("Searching for colony locations automatically.")
            print("Checking "+str(inp.updates)+" (quasi-random) candidate grid positions in first phase of grid location")
        if inp.greenlab:
            print("Removing images identified as containing GreenLab lids.")
        if inp.lc:
            if inp.cut:
                print("Cutting cell signal from first image to create pseudo-empty plate (for lighting correction).")
            else:
                print("Using first plate (without segmenting) as best estimate of pseudo-empty plate (for lighting correction).")
        if fixedThresh==-99 and not inp.edgemask:
            print("Image segmentation by automatic thresholding.")
        elif not inp.edgemask:
            print("Images will be segmented using fixed threshold: "+str(fixedThresh)+".")
        else:
            print("Images will be segemented by intensity gradient and morphology instead of by thresholding.")
        if fdict is not None and os.path.exists(fdict):
            print("Preparing to load barcodes from "+fdict+".")
    res={'lc':inp.lc,'fixedThresh':fixedThresh,'plots':inp.plots,'initpos':inp.initpos,'fdict':fdict,'fdir':fdir,'nrow':nrow,'ncol':ncol,'cut':cut,'verbose':verbose,'diffims':diffIms,'updates':inp.updates,'endpoint':inp.endpoint,'edgemask':inp.edgemask,'greenlab':inp.greenlab}
    return(res)

def locateJSON(scrID,dirHTS='.',verbose=False):
    exptType=scrID[0:-4]
    fdict=os.path.join(dirHTS,exptType+"_EXPERIMENTS",scrID,"AUXILIARY",scrID+"_C2.json")
    return(fdict)

def prepareTimecourse(barcdict,verbose=False,chklid=lambda x:False):
    '''In timecourse mode, prepares "next" batch of images for analysis from dictionary of image names (unique image barcodes are dictionary keys).'''
    BARCs=sorted(barcdict)
    BARCODE=BARCs[0]
    imdir=os.path.dirname(barcdict[BARCODE][0])
    IMs=[b for b in barcdict[BARCODE] if not chklid(b)]
    LATESTIMAGE=IMs[0]
    EARLIESTIMAGE=IMs[-1]
    imRoot=EARLIESTIMAGE.split(".")[0]
    if verbose:
        print("Analysing images labelled with the barcode "+BARCODE+" in "+imdir)
        print("Earliest image: "+EARLIESTIMAGE)
        print("Latest image: "+LATESTIMAGE)
    return((BARCODE,imdir,LATESTIMAGE,EARLIESTIMAGE,imRoot))

def loadLocationGuesses(IMAGE,InsData):
    '''Set up initial guesses for location of (centres of) spots on image by parsing data from Colonyzer.txt'''
    # If we have ColonyzerParametryzer output for this filename, use it for initial culture location estimates
    if os.path.basename(IMAGE) in InsData:
        (candx,candy,dx,dy)=c2.SetUp(InsData[os.path.basename(IMAGE)])
    # If there are multiple calibrations available, choose the best one based on date of image capture
    elif any(isinstance(el, list) for el in InsData['default']):
        imname=os.path.basename(IMAGE).split(".")[0]
        imdate=imname[-19:-9]
        (candx,candy,dx,dy)=c2.SetUp(InsData['default'],imdate)
    else:
        (candx,candy,dx,dy)=c2.SetUp(InsData['default'])
    return((candx,candy,dx,dy))

def cutEdgesFromMask(mask,locations,dx,dy):
    '''Mask for identifying culture areas (edge detection). Set all pixels outside culture grid to background, to aid binary filling later.'''
    mask[0:min(locations.y-dy/2),:]=False
    mask[max(locations.y+dy/2):mask.shape[0],:]=False
    mask[:,0:min(locations.x-dx/2)]=False
    mask[:,max(locations.x+dx/2):mask.shape[1]]=False
    return(mask)

def edgeFill(arr,locations,dx,dy,cutoff=0.8):
    edgeN=c2.getEdges(arr,cutoff=cutoff)
    dilN=ndimage.morphology.binary_dilation(edgeN,iterations=2)
    erodeN=ndimage.morphology.binary_erosion(dilN,iterations=1)
    dil2N=ndimage.morphology.binary_dilation(dilN,iterations=3)
    
    fillN=ndimage.morphology.binary_fill_holes(cutEdgesFromMask(dil2N,locations,dx,dy))
    maskN=ndimage.morphology.binary_erosion(fillN,iterations=7)
    return(maskN)

def edgeFill2(arr,cutoff=0.8):
    edgeN=c2.getEdges(arr,cutoff=cutoff)
    er1=ndimage.morphology.binary_erosion(edgeN,iterations=2)
    di1=ndimage.morphology.binary_dilation(er1,iterations=4)
    edgemap=np.logical_and(di1,edgeN)
    fillN=ndimage.morphology.binary_fill_holes(edgemap)
    return(fillN)

def main(inp=""):
    print("Colonyzer "+c2.__version__)
    
    var=buildVars(inp=inp)
    correction,fixedThresh,plots,initpos,fdict,fdir,nrow,ncol,cut,verbose,diffIms,updates,endpoint,edgemask,greenlab=(var["lc"],var["fixedThresh"],var["plots"],var["initpos"],var["fdict"],var["fdir"],var["nrow"],var["ncol"],var["cut"],var["verbose"],var["diffims"],var["updates"],var["endpoint"],var["edgemask"],var["greenlab"])
    barcdict=checkImages(fdir,fdict,verbose=verbose)
    if greenlab:
        barcdict={x:barcdict[x] for x in barcdict.keys() if not c2.checkAnalysisStarted(barcdict[x][-1])}
    rept=c2.setupDirectories(barcdict,verbose=verbose)

    if greenlab:
        posfiles=[resource_filename(Requirement.parse("colonyzer2"),os.path.join("data",f)) for f in ["GreenLabLid.png","CornerLid.png","BottomRightLid.png"]]
        negfiles=[resource_filename(Requirement.parse("colonyzer2"),os.path.join("data",f)) for f in ["GreenLabNoLid.png","CornerNoLid.png","BottomRightNoLid.png"]]
        frects=[[0.15,0.85,0.5,1.0],[0.0,0.8,0.125,1.0],[0.85,0.8,1.0,1.0]]
        pdistThresh,ndistThresh=50.3,53.333
        checkLid=c2.makeLidTest(posfiles,negfiles,frects,pdistThresh,ndistThresh,False)
    else:
        def checkLid(fname): False

    start=time.time()

    while len(barcdict)>0:

        BARCODE,imdir,LATESTIMAGE,EARLIESTIMAGE,imRoot=prepareTimecourse(barcdict,verbose=verbose,chklid=checkLid)  

        if plots:
            pdf=PdfPages(os.path.join(os.path.dirname(EARLIESTIMAGE),"Output_Reports",os.path.basename(EARLIESTIMAGE).split(".")[0]+".pdf"))
        else:
            pdf=None
        
        # Create empty file to indicate that barcode is currently being analysed, to allow parallel analysis (lock files)
        tmp=open(os.path.join(os.path.dirname(EARLIESTIMAGE),"Output_Data",os.path.basename(EARLIESTIMAGE).split(".")[0]+".out"),"w").close()

        # Get latest image for thresholding and detecting culture locations
        imN,arrN=c2.openImage(LATESTIMAGE)
        # Get earliest image for lighting gradient correction
        if (LATESTIMAGE==EARLIESTIMAGE) or endpoint:
            im0,arr0=imN,arrN
            arrloc=arrN
        else:
            im0,arr0=c2.openImage(EARLIESTIMAGE)
            arrloc=arrN-arr0
            #arrloc=np.maximum(0,arrloc)
        if initpos:
            InsData=c2.readInstructions(os.path.dirname(LATESTIMAGE),searchUpStream=True)
            # Load initial guesses from Colonyzer.txt file
            (candx,candy,dx,dy)=loadLocationGuesses(LATESTIMAGE,InsData)
            corner=[0,0]; com=[0,0]; guess=[0,0]
            # NOTE: assumes that grid returned by loadLocationGuesses is parallel to image edges
            ny=nrow=len(np.unique(candy)) 
            nx=ncol=len(np.unique(candx))
        else:
            # Automatically generate guesses for gridded array locations
            (candx,candy,dx,dy,corner,com,guess)=c2.estimateLocations(arrloc,ncol,nrow,showPlt=plots,pdf=pdf,glob=False,verbose=verbose,nsol=updates)

        # Update guesses and initialise locations data frame
        locationsN=c2.locateCultures([int(round(cx-dx/2.0)) for cx in candx],[int(round(cy-dy/2.0)) for cy in candy],dx,dy,arrloc,ncol,nrow,update=True)

        mask=edgeFill2(arrN,0.9)
        c2.showIm(mask)
        grd=mask.copy()
        grd[:,:]=False
        grd[min(locationsN.y-dy/2):max(locationsN.y+dy/2),min(locationsN.x-dx/2):max(locationsN.x+dx/2)]=True
        spots=np.logical_and(grd,mask)
        agar=np.logical_and(grd,~mask)
        ave0=np.mean(arr0[agar])
        thresh=-99
        
        if correction:
            if cut:
                pseudoempty = np.array(np.round(arr0.copy()),dtype=np.uint8)
                bytemask = np.asarray(spots*255, dtype=np.uint8)
                filled = cv2.inpaint(pseudoempty,bytemask,3,cv2.INPAINT_NS)
                blurred = cv2.GaussianBlur(filled,(21,21),0)
                pseudoempty[spots] = blurred[spots]
            else:
                pseudoempty=arr0
        else:
            pseudoempty=0

        for FILENAME in barcdict[BARCODE]:

            startim=time.time()
            
            im,arr=c2.openImage(FILENAME)
            # Local updates on individual images to allow for slight movement of plate
            #locationsN=c2.locateCultures([int(round(cx-dx/2.0)) for cx in candx],[int(round(cy-dy/2.0)) for cy in candy],dx,dy,arr,ncol,nrow,update=True)

            ave=np.mean(arr[agar])
            # Correct for minor lighting differences between images
            arr=np.maximum(0,np.minimum(255,arr-(ave-ave0)))
            # Subtract background (corrects for lighting differences within plate/image as well as making agar intensity correspond to zero signal)
            arr=np.maximum(arr-pseudoempty,0)
            
            # Measure culture phenotypes
            locations=c2.measureSizeAndColour(locationsN,arr,im,spots,0,BARCODE,FILENAME[0:-4])

            # Write results to file
            locations.to_csv(os.path.join(os.path.dirname(FILENAME),"Output_Data",os.path.basename(FILENAME).split(".")[0]+".out"),"\t",index=False,engine='python')
            dataf=c2.saveColonyzer(os.path.join(os.path.dirname(FILENAME),"Output_Data",os.path.basename(FILENAME).split(".")[0]+".dat"),locations,thresh,dx,dy)

            # Visual check of culture locations
            imthresh=c2.threshPreview(locations,mask,None)

            r=5
            draw=ImageDraw.Draw(imthresh)
            draw.ellipse((com[1]-r,com[0]-r,com[1]+r,com[0]+r),fill=(255,0,0))
            draw.ellipse((corner[1]-r,corner[0]-r,corner[1]+r,corner[0]+r),fill=(255,0,0))
            draw.ellipse((guess[1]-r,guess[0]-r,guess[1]+r,guess[0]+r),fill=(255,0,0))
            draw.ellipse((candx[0]-r,candy[0]-r,candx[0]+r,candy[0]+r),fill=(255,0,0))
            imthresh.save(os.path.join(os.path.dirname(FILENAME),"Output_Images",os.path.basename(FILENAME).split(".")[0]+".png"))

            # Get ready for next image
            if verbose: print("Finished {0} in {1:.2f}s".format(os.path.basename(FILENAME),time.time()-startim))

        # Get ready for next image
        if verbose: print("Finished {0} in {1:.2f}s".format(os.path.basename(BARCODE),time.time()-start))

        barcdict={x:barcdict[x] for x in barcdict.keys() if not c2.checkAnalysisStarted(barcdict[x][-1])}
        if plots:
            pdf.close()

    print("No more barcodes to analyse... I'm done.")

if __name__ == '__main__':
    main()
