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

    def pycosmic_multi():

        # 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 = outputs.out_sof.files

        masterbiases={}
        badpixmaps={}
        for file in files:
            if file.category == 'MASTER_BIAS' :
                for p in file.purposes :
                    masterbiases[p]=file
            if file.category == 'BAD_PIXEL_MAP' :
                for p in file.purposes :
                    badpixmaps[p]=file

        #Define PyCosmic parameters:
        args=lambda: None
        args.rdnoise="HIERARCH ESO DET OUT1 RON"
        args.siglim=5.0
        args.fwhm=1.7 #2.5 # 2.0
        args.rlim=0.5 # 1.2
        args.iter=5 #5
        #args.iter=1 #5
        args.replacebox=[5,5]
        args.radius=0
        args.gain="HIERARCH ESO DET OUT1 CONAD"
        args.parallel=True
        args.verbose=True

        outlist=list()
        for file in files:
            #select a file category:
            if file.category in ['SCIENCE', 'STD' ] :
                purpose=file.purposes[0]
                if purpose in masterbiases :
                    try:
                        bHDU=pyfits.open(masterbiases[purpose].name)
                    except:
                        try:
                            logging.error('Could not open MASTER_BIAS %s' %(masterbiases[0].name))
                        except:
                            pass
                        sys.exit(1)
                if purpose in badpixmaps :
                    try:
                        bpmHDU=pyfits.open(badpixmaps[purpose].name)
                    except:
                        try:
                            logging.error('Could not open BAD_PIXEL_MAP %s' %(badpixmaps[0].name))
                        except:
                            pass
                        sys.exit(1)
                # define file names:
                raw_file=file.name
                bs_file=parameters["--products-dir"] + "/" + \
                       raw_file.replace(".fits", "_bs.fits").split("/")[-1]
                clean_file=parameters["--products-dir"] + "/" + \
                       raw_file.replace(".fits", "_PC_clean.fits").split("/")[-1]
                mask_file=parameters["--products-dir"] + "/" + \
                       raw_file.replace(".fits", "_PC_mask.fits").split("/")[-1]
                file.name=clean_file

                bs=pyfits.open(raw_file)
                if purpose in masterbiases :
                    bs[0].data=bs[0].data-bHDU[0].data
                    #bs[0].header['BZERO']=0.
                    if 'BZERO' in  bs[0].header :
                        del bs[0].header['BZERO']
                if purpose in badpixmaps :
                    bs.append(bpmHDU[0])
                    bs[-1].header['EXTNAME']='BADPIX'

                try:
                    bs.writeto(bs_file, output_verify='silentfix+ignore', overwrite=True)
                except:
                    bs.writeto(bs_file, output_verify='silentfix+ignore', clobber=True)

                # detCos
                #Detects and removes cosmics from astronomical images based on Laplacian edge 
                #detection scheme combined with a PSF convolution approach (Husemann  et al. in prep.).
           
                #IMPORTANT: 
                #The image and the readout noise are assumed to be in units of electrons.
                #The image also needs to be BIAS subtracted!
                #The gain can be entered to convert the image from ADUs to electros, when this is down already set gain=1.0 as the default.
                # clean cosmics:
                PyCosmic.detCos(bs_file, mask_file, clean_file, args.rdnoise, gain=args.gain, sigma_det=args.siglim,  rlim=args.rlim, iter=args.iter, fwhm_gauss=args.fwhm, replace_box=args.replacebox,    increase_radius=args.radius, parallel=args.parallel, verbose=args.verbose)

                #...the output of PyCosmic is multiplied by the gain, divide it out again...
                clean=pyfits.open(clean_file, 'update')

                if purpose in masterbiases :
                    clean[0].data=clean[0].data/clean[0].header[args.gain]+bHDU[0].data
                else :
                    clean[0].data=clean[0].data/clean[0].header[args.gain]

                # update header
                clean[0].header.append(("HIERARCH ESO PRO CMETHOD", "PYCOSMIC", "method for bad pixel cleaning"))
                clean[0].header.append(("HIERARCH ESO DRS CMETHOD", "PYCOSMIC", "method for bad pixel cleaning"))
                clean[0].header.append(("CMETHOD", "PYCOSMIC", "method for bad pixel cleaning"))
                clean[0].header.append(("NFILES", len(files), "number of CR cleaned files"))
                clean[0].header.append(("CRSIGLIM", args.siglim, "PyCosmic signifiance threshold"))
                clean[0].header.append(("CRRLIM", args.rlim, "PyCosmic contrast limit"))
                clean[0].header.append(("CRITER", args.iter, "PyCosmic number of iterations"))
                clean[0].header.append(("CRGFWHM", args.fwhm, "PyCosmic FWHM of convolution kernel"))
                clean[0].header.append(("CRRBOX", str(args.replacebox), "PyCosmic box to computer cleaned value"))
                clean[0].header.append(("CRRADIUS", args.radius, "PyCosmic number of neighboring pixels"))
                clean[0].header.append(('CRRON',clean[0].header['HIERARCH ESO DET OUT1 RON'],'Readout noise'))
                clean[0].header.append(('CRGAIN',clean[0].header['HIERARCH ESO DET OUT1 CONAD'],'Gain'))
                if purpose in masterbiases :
                    clean[0].header.append(('CRMBIAS',os.path.basename(masterbiases[purpose].name)[0:80-len('CRMBIAS')-5],'MASTER_BIAS used'))
                if purpose in badpixmaps :
                    clean[0].header.append(('CRBPM',os.path.basename(badpixmaps[purpose].name)[0:80-len('CRBPM')-5],'BAD_PIX_MAP used'))
                clean[0].header.append(('HIERARCH ESO DRS PYCOSMIC',True,'AstroScrappy Cosmic Ray cleaning applied'))
                if purpose in masterbiases :
                    clean[0].header.append(('HIERARCH ESO DRS PYCOSMIC BIASSUB',True,'MASTER_BIAS subtracted?'))
                    clean[0].header.append(('HIERARCH ESO DRS PYCOSMIC BIAS NAME',os.path.basename(masterbiases[purpose].name)[0:80-len('HIERARCH ESO PYCOSMIC BIAS NAME')-5],'MASTER BIAS used'))
                clean[0].header.append(('HIERARCH ESO DRS PYCOSMIC RON',clean[0].header['HIERARCH ESO DET OUT1 RON'],'Readout noise'))
                clean[0].header.append(('HIERARCH ESO DRS PYCOSMIC GAIN',clean[0].header['HIERARCH ESO DET OUT1 CONAD'],'Gain'))
                clean.flush()
                clean.close()

                # append remaining extensions from raw files unchanged:
                raw=pyfits.open(raw_file)
                for iext in range(1,len(raw)):
                    if import_astropy_io :
                        pyfits.append(clean_file, raw[iext].data, header=raw[iext].header, checksum=True)
                    else :
                        pyfits.append(clean_file, raw[iext].data, header=raw[iext].header)
                raw.close()
                #clean=pyfits.open(clean_file, 'update')
                #print(numpy.min(clean[0].data[:]),numpy.max(clean[0].data[:]))
                #clean.close()

                # update purpose
                outlist.append((reflex.FitsFile(file.name,file.category,None,[str("PC_crclean/")+file.purposes[0]])))

        #send output to EsoReflex:
        newsof = reflex.SetOfFiles(inputs.in_sof.datasetName,outlist)
        outputs.out_sof = newsof
        #parser.write_outputs()
        #for file in outlist:
            #print("output:", file.name, file.category, file.purposes)


except ImportError:
    import_sucess = False
    print("Error importing one or more of the modules astropy.io/pyfits, PyCosmic")

except :
    raise

if __name__ == '__main__':
  
    # Hacked from...
    # example script to process input files and produce one output file for
    # each processed input.

    #send output to EsoReflex:
    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 : #and inputs.enable == "true" :
        pycosmic_multi()
    else :
        # write out an empty sof
        outputs.out_sof = reflex.SetOfFiles(inputs.in_sof.datasetName,list())

    #send output to EsoReflex:
    parser.write_outputs()

    sys.exit()
