# ========================================================================================================
# import the needed modules
try:
    import reflex 
    import sys
    import os
    try :
        from astropy.io import fits as fits
    except:
        import pyfits as fits
    import numpy
    import_sucess = True

    def median_multi( enable=False ):

        # get Kepler parameters:
        parameters= {}
        for p in sys.argv[1:]:
            parameters[p.split("=")[0]]= p.split("=")[1]

        #Define one output file for every input file:
        #outputs.out_sof = inputs.in_sof
        files = inputs.in_sof.files

        scifiles = []
        for file in files :
            if file.category == 'SCIENCE' or file.category == 'STD' :
                scifiles.append(file)

        outlist=list()
        if len(scifiles) > 0 :
            cube=[]
            median_files=[]
            oldpurpose=scifiles[0].purposes[0]
            for ifile,file in enumerate(scifiles):
                # define file names:
                raw_file=file.name
                raw=fits.open(raw_file)
                if enable and ifile > 1:
                    cube.append(raw[0].data)
                    median_image=numpy.median(numpy.array(cube), axis=0)

                    firstraw=fits.open(scifiles[0].name)
                    firstraw[0].data=median_image
                    if 'BZERO' in firstraw[0].header :
                        del firstraw[0].header['BZERO']

                    # save as new file:
                    if ifile == len(scifiles)-1:
                        median_file= parameters["--products-dir"] + "/" + \
	                       (scifiles[0].name).replace(".fits", "_median.fits").split("/")[-1]
                    else:
                        median_file= parameters["--products-dir"] + "/" + \
	                       (scifiles[0].name).replace(".fits", "_median"+str(ifile+1)+".fits").split("/")[-1]
                    median_files.append(median_file)

                    # update header:
                    for i in range(ifile):
                        filename=(scifiles[i].name).split("/")[-1]
                        firstraw[0].header.append(("IFILE"+str(i), filename, "input file for median"))
                    firstraw[0].header.append(("NFILES", ifile+1, "number of input files"))
                    if ifile == len(scifiles)-1:
                        firstraw[0].header.append(("CMETHOD", "MEDIAN", "method for bad pixel cleaning"))
                        firstraw[0].header.append(("HIERARCH ESO PRO CMETHOD", "MEDIAN", "method for bad pixel cleaning"))
                        firstraw[0].header.append(("HIERARCH ESO DRS CMETHOD", "MEDIAN", "method for bad pixel cleaning"))
                    else:
                        firstraw[0].header.append(("CMETHOD", "SUBMED_%d" %(ifile+1), "method for bad pixel cleaning"))
                        firstraw[0].header.append(("HIERARCH ESO PRO CMETHOD", "SUBMED_%d" %(ifile+1), "method for bad pixel cleaning"))
                        firstraw[0].header.append(("HIERARCH ESO DRS CMETHOD", "SUBMED_%d" %(ifile+1), "method for bad pixel cleaning"))

                    # write fits file:
                    try:
                        firstraw.writeto(median_file, output_verify='silentfix+ignore', checksum=True, overwrite=True)
                    except:
                        firstraw.writeto(median_file, output_verify='silentfix+ignore', checksum=True, clobber=True)
                    firstraw.close()

                '''
                # Don't do this now that we have implemented optional keys in the renamed filename
                # And write out a copy of the individual raw files
                raw[0].header.append(("NFILES", 1, "number of input files"))
                raw[0].header.append(("CMETHOD", "NONE", "method for bad pixel cleaning"))
                raw[0].header.append(("HIERARCH ESO PRO CMETHOD", "NONE", "method for bad pixel cleaning"))
                raw[0].header.append(("HIERARCH ESO DRS CMETHOD", "NONE", "method for bad pixel cleaning"))
                try:
                    raw.writeto(parameters["--products-dir"] + "/" + os.path.basename(file.name), overwrite=True)
                except:
                    raw.writeto(parameters["--products-dir"] + "/" + os.path.basename(file.name), clobber=True)
                '''
                # create output SOF:
                # first, lets put in the individual files with a new purpose:
                #outlist.append((reflex.FitsFile(parameters["--products-dir"] + "/" + os.path.basename(file.name),file.category,None,[str("orig/")+oldpurpose])))
                outlist.append((reflex.FitsFile(file.name,file.category,None,[str("orig/")+oldpurpose])))
                raw.close()

            # now add the new median files:
            if enable :
                if False :
                    # Add all the SUBMED versions...
                    for imed, median_file in enumerate(median_files):
                        outlist.append((reflex.FitsFile(median_file,scifiles[0].category,None,[str("med"+str(imed+3)+"/") + oldpurpose])))
                elif len(median_files) > 0 :
                    # Add only the final full median...
                    outlist.append((reflex.FitsFile(median_files[-1],scifiles[0].category,None,[str("med/") + oldpurpose])))

            #for file in outlist:
            #    print("output:", file.name, file.category, file.purposes)

        #send output to EsoReflex:
        newsof = reflex.SetOfFiles(inputs.in_sof.datasetName,outlist)
        outputs.out_sof = newsof


except ImportError:
    import_sucess = False
    print("Error importing one or more of the modules pyfits")

if __name__ == '__main__':

    # Hacked from...
    # example script to process input files and produce one output file for
    # each processed input.
    parser = reflex.ReflexIOParser()

    #Define ports:
    parser.add_option("-i", "--in_sof", dest="in_sof")
    parser.add_option("-e", "--enable", dest="enable", default='false')
    parser.add_output("-o", "--out_sof", dest="out_sof")

    inputs  = parser.get_inputs()
    outputs = parser.get_outputs()
  
    #Check if import failed or not
    if import_sucess :
        median_multi( inputs.enable.lower() == 'true' )
    else :
        outlist=list()
        for file in inputs.in_sof.files:
            if file.category == 'SCIENCE' or file.category == 'STD' :
                outlist.append((reflex.FitsFile(file.name,file.category,file.checksum,file.purposes)))

        outputs.out_sof = reflex.SetOfFiles(inputs.in_sof.datasetName,outlist)

    #send output to EsoReflex:
    parser.write_outputs()

    sys.exit()
