# 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
            ## For the moment using the Bad Pixel map trying causes PyCosmic to crash, so don't...
            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 == 'SCIENCE' or file.category == 'STD' :
                purpose=file.purposes[0]
                if purpose in masterbiases :
                    try:
                        ## For the moment we are just using the first MASTER_BIAS, but not sure that is correct...
                        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:
                        ## For the moment we are just using the first MASTER_BIAS, but not sure that is correct...
                        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, overwrite=True)
                except:
                    bs.writeto(bs_file, 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')

                clean[0].data=clean[0].data/clean[0].header[args.gain]+bHDU[0].data

                # update header
                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.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,[unicode("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"

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_output("-o", "--out_sof", dest="out_sof")

    inputs  = parser.get_inputs()
    outputs = parser.get_outputs()

    #Check if import failed or not
    if import_sucess :
        pycosmic_multi()
    else :
        outputs.out_sof = reflex.SetOfFiles(inputs.in_sof.datasetName,list())

    #send output to EsoReflex:
    parser.write_outputs()

    sys.exit()
