from __future__ import print_function
##CUSP UO 2016
__author__ = "fbb"


import os
import sys
import glob
import subprocess
import optparse

import numpy as np
import pylab as pl
from findImageSize import findsize

def runit((arg, options)):
    filepattern = arg
    fnameroot = filepattern.split('/')[-1]
    impath = os.getenv("UIdata") + filepattern
    print ("\n\nUsing image path: %s\n\n"%impath)    
    flist = sorted(glob.glob(impath+"*.raw"))
    ##print (flist)
    if options.stack:
        print ("(There is a stack", options.stack,")")
        stack = np.load(options.stack)
        imsize  = findsize(stack,
                           filepattern=options.stack.replace('.npy','.txt'))
    else:
        print ("Please pass a stack with option --stack")
    
    print ("Image size: ", imsize)
        
    nmax = min(options.nmax, len(flist)-options.skipfiles)
    print ("Number of timestamps (files): %d"%(nmax))

    if nmax<30: return 0
    flist = flist[options.skipfiles:nmax+options.skipfiles]    
    imgmask = np.load(options.stack.replace(".npy","_mask.npy"))
    if not options.lmax:
        lmax = (imgmask.sum(axis=2).sum(axis=1)>0).sum()
    else: lmax  = options.lmax
    print (lmax)
    outdir = '/'.join(filepattern.split('/')[:-1])+'/N%04dW%04dS%04d'%(nmax,lmax,options.skipfiles)
    
    if not os.path.isdir(outdir):
        subprocess.Popen('mkdir -p %s '%outdir, shell=True)
        #os.system('mkdir -p %s '%outdir)
    print ("Output directory: ",
           '/'.join(filepattern.split('/')[:-1]),  outdir)

    if options.fft:
        #bsoutfile = outdir + "/" + fnameroot + "_bs_fft.npy"
        #coordsoutfile = outdir + "/" + fnameroot + "_coords_fft.npy"
        goodcoordsoutfile = outdir + "/" + fnameroot + "_goodcoords_fft.npy"  
        #kmresultfile = outdir + "/" + fnameroot + "_kmresult_fft.pkl"

        
    else:
        #bsoutfile = outdir + "/" + fnameroot + "_bs.npy"
        #coordsoutfile = outdir + "/" + fnameroot + "_coords.npy"
        goodcoordsoutfile = outdir + "/" + fnameroot + "_goodcoords.npy"
        #kmresultfile = outdir + "/" + fnameroot + "_kmresult.pkl"

    print ("")

    print ("must read: ", goodcoordsoutfile)

    pl.figure()
    img = np.load(options.stack)
    imgm = imgmask.sum(axis=0)
    pl.imshow(img)
    pl.savefig(filepattern+"_mask.stack.png")
    imgma = np.zeros((img.shape[0],imgm.shape[1],4))
    imgma[imgm==1]=[1,1,1,1]
    pl.imshow(imgma)
    pl.savefig(filepattern+"_mask.img.png")
    
    print ("saved figure mask", filepattern+"_mask.img.png")
    gw = np.load(outdir+"/"+filepattern.split('/')[-1]+"_goodcoords.npy")
    imgmagood = np.zeros((img.shape[0],imgm.shape[1],4))

    for g in gw.T:
        if not np.isnan(g[2]) and not np.isnan(g[3]):
            imgmagood[:,:,0] = imgmagood[:,:,0]+(imgmask[imgmask[:,int(g[3]),int(g[2])]>0]).squeeze()
            imgmagood[:,:,3] = imgmagood[:,:,3]+(imgmask[imgmask[:,int(g[3]),int(g[2])]>0]).squeeze()
    pl.imshow(img)
    pl.imshow(imgma)
    pl.imshow(imgmagood)
    pl.savefig(filepattern+"_goodmask.img.png")
    
    


if __name__=='__main__':

    parser = optparse.OptionParser(usage="getallcv.py 'filepattern' ",
                                   conflict_handler="resolve")
    parser.add_option('--nmax', default=100, type="int",
                      help='number of images to process (i.e. timestamps)')
    parser.add_option('--lmax', default=None, type="int",
                      help='number of lights')
    parser.add_option('--showme', default=False, action="store_true",
                      help='show plots')
    parser.add_option('--coordfile', default=None, type="str",
                      help='coordinates python array (generated by windowFinder.py)')
    parser.add_option('--stack', default=None, type="str",
                      help='stack python array')
    parser.add_option('--skipfiles', default=150, type="int",
                      help="number of files to skip at the beginning")
    parser.add_option('--fft', default=False, action="store_true",
                      help='custer in fourier space')
    
    options,  args = parser.parse_args()
    
    print ("options", options)
    print ("args", args, args[0])
    if len(args) < 1:
        sys.argv.append('--help')
        options,  args = parser.parse_args()
           
        sys.exit(0)
    for arg in args:
        runit((arg, options))
    
