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

    def astroscrappy_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

        outlist=list()
        for file in files:
            #select a file category:
            if file.category in ['SCIENCE', 'STD' ] :

                purpose=file.purposes[0]
                try: 
                    HDU=pyfits.open(file.name)
                except:
                    try:
                        logging.error('Could not open %s' %(file.name))
                    except:
                        pass
                    sys.exit(1)
                if purpose in masterbiases :
                    try:
                        bHDU=pyfits.open(masterbiases[purpose].name)
                    except:
                        try:
                            logging.error('Could not open MASTER_BIAS %s' %(masterbiases[purpose].name))
                        except:
                            pass
                        sys.exit(1)

                if not 'CRMASK' in HDU :
                    HDU.append(HDU[0])
                    HDU[-1].header['EXTNAME']='CRMASK'
                    HDU['CRMASK'].header['HIERARCH ESO PRO CATG']='CRMASK'
                    

                if purpose in masterbiases :
                    mask, HDU[0].data = astroscrappy.detect_cosmics(
                        HDU[0].data-bHDU[0].data, 
                        gain=HDU[0].header['HIERARCH ESO DET OUT1 CONAD'], 
                        readnoise=HDU[0].header['HIERARCH ESO DET OUT1 RON']
                      )
                else :
                    mask, HDU[0].data = astroscrappy.detect_cosmics(
                        HDU[0].data,
                        gain=HDU[0].header['HIERARCH ESO DET OUT1 CONAD'], 
                        readnoise=HDU[0].header['HIERARCH ESO DET OUT1 RON']
                      )

                #HDU[0].header['BZERO']=0.
                if 'BZERO' in  HDU[0].header :
                    del HDU[0].header['BZERO']
                HDU[0].data=HDU[0].data/HDU[0].header['HIERARCH ESO DET OUT1 CONAD']
                if purpose in masterbiases :
                    HDU[0].data+=bHDU[0].data
                HDU['CRMASK'].data=mask.astype(int)
                HDU['CRMASK'].header['BZERO']=0.
                HDU[0].header.append(("NFILES", 1, "number of input files"))
                HDU[0].header.append(("CMETHOD", "ASTROSCRAPPY", "method for bad pixel cleaning"))
                HDU[0].header.append(("HIERARCH ESO DRS CMETHOD", "ASTROSCRAPPY", "method for bad pixel cleaning"))
                HDU[0].header.append(('HIERARCH ESO DRS ASCRAP',True,'AstroScrappy Cosmic Ray cleaning applied'))
                if purpose in masterbiases :
                    HDU[0].header.append(('HIERARCH ESO DRS ASCRAP BIASSUB',True,'MASTER_BIAS subtracted?'))
                    HDU[0].header.append(('HIERARCH ESO DRS ASCRAP BIAS NAME',os.path.basename(masterbiases[purpose].name)[0:80-len('HIERARCH ESO ASCRAP BIAS NAME')-5],'MASTER BIAS used'))
                HDU[0].header.append(('HIERARCH ESO DRS ASCRAP RON',HDU[0].header['HIERARCH ESO DET OUT1 RON'],'Readout noise'))
                HDU[0].header.append(('HIERARCH ESO DRS ASCRAP GAIN',HDU[0].header['HIERARCH ESO DET OUT1 CONAD'],'Gain'))

                # define file names:
                clean_file= parameters["--products-dir"] + "/" + \
                       file.name.replace(".fits", "_AS_clean.fits").split("/")[-1]
                file.name=clean_file
                try:
                    HDU.writeto(clean_file, outpur_verify='silentfix+ignore', overwrite=True)
                except:
                    HDU.writeto(clean_file, output_verify='silentfix+ignore', clobber=True)
                HDU.close()
                if purpose in masterbiases :
                    bHDU.close()

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

        #send output to EsoReflex:
        newsof = reflex.SetOfFiles(inputs.in_sof.datasetName,outlist)
        outputs.out_sof = newsof
        #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, astroscrappy")

except :
    raise

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

    parser.write_outputs()

    sys.exit()
