# ========================================================================================================
# 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 numpy as np
    from scipy import interpolate
    import copy

    import_sucess = True

    def pycosmic_multi(nan_corr_method='1D'):

        # 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]
                # Fix any NaNs...
                number_of_NaNs=0
                NaN_correction_method='None'
                if np.any(np.isnan(clean[0].data)) :
                    _o=pyfits.open(raw_file)
                    _d=np.shape(clean[0].data)
                    _x=np.tile(np.arange(_d[1]),_d[0])
                    _y=np.repeat(np.arange(_d[0]),_d[1])
                    nans, x = nan_helper(np.ravel(clean[0].data))
                    number_of_NaNs=len(_x[nans])
                    NAN_INTERP_METHOD=os.environ.get('GIRAF_PYCOSMIC_NAN_INTERP_METHOD',nan_corr_method).lower()
                    print('Control the NAN_INTERP_METHOD by setting GIRAF_PYCOSMIC_NAN_INTERP_METHOD=["1D", "2D", "orig_pix_value"].')
                    print('Default value is "1D".')
                    print("Note: ENV-VAR method doesn't work in esoreflex environment if esoreflex.inherit-environment=FALSE.")
                    print("Instead set it in the workflow.")
                    if NAN_INTERP_METHOD in ['orig_pix_value',] :
                        ## Replace the NaN with the original pixel value...
                        print('Fixing %d NaNs by replacing with original pixel value...' %(number_of_NaNs))
                        NaN_correction_method='OPIXVAL'
                        for [i,j] in np.transpose(np.reshape(list(_y[nans])+list(_x[nans]),[2,len(_x[nans])])) :
                            clean[0].data[i,j]=_o[0].data[i,j]
                    elif NAN_INTERP_METHOD in ['2d',] :
                        ## 2-D interpolation...
                        print('Fixing %d NaNs with 2-D interpolation...' %(number_of_NaNs))
                        NaN_correction_method='2D-INTERP'
                        for [i,j] in np.transpose(np.reshape(list(_y[nans])+list(_x[nans]),[2,len(_x[nans])])) :
                            _npx=7
                            _npy=7
                            __x=np.tile(np.arange(_npx)-_npx//2,_npy)
                            __y=np.repeat(np.arange(_npy)-_npy//2,_npx)
                            rclean=np.ravel(cd[
                                np.max([0,(i-_npx//2)]):np.min([(i+_npx//2+1),_d[0]-1]),
                                np.max([0,(j-_npy//2)]):np.min([(j+_npy//2+1),_d[1]-1])
                            ])
                            ## Here we assume that the number of NaNs in the 7x7 box is sufficiently small that we can get
                            ## a good interpolation from the remaining NaNs
                            ## The one example looked at so far (GIRAF.2007-04-09T06:52:05.052) had only a single NaN
                            nans, x = nan_helper(rclean)
                            rclean[nans]=0.
                            _f=interpolate.interp2d(__x[x(~nans)],__y[x(~nans)],rclean[x(~nans)], kind='linear')
                            clean[0].data[i,j]=_f(0,0)[0]
                    else :
                        ## 1-D interpolation
                        # Interpolate the values of NaN pixels from the neighbouring pixels of the same column...
                        print('Fixing %d NaNs with 1-D interpolation...' %(number_of_NaNs))
                        NaN_correction_method='1D-INTERP'
                        nans, x = nan_helper(np.sum(clean[0].data,axis=0))
                        for i in x(nans) :
                            nnans, xx = nan_helper(clean[0].data[:,i])
                            clean[0].data[:,i][nnans]= np.interp(
                                xx(nnans), xx(~nnans), clean[0].data[:,i][~nnans]
                            )
                    '''
                    clean.append(pyfits.ImageHDU(header=clean[0].header,data=_o[0].data-clean[0].data))
                    clean[-1].header['HIERARCH ESO PRO CATG']="NAN_DIFF"
                    '''
                # 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(("HIERARCH ESO DRS NANS NUM", number_of_NaNs, "Number of NaNs detected in cleaned frame"))
                clean[0].header.append(("HIERARCH ESO DRS NANS METHOD", NaN_correction_method, "Method used to correct NaNs"))
                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

# ========================================================================================================
## From https://stackoverflow.com/questions/6518811/interpolate-nan-values-in-a-numpy-array
def nan_helper(y):
    """Helper to handle indices and logical indices of NaNs.

    Input:
        - y, 1d numpy array with possible NaNs
    Output:
        - nans, logical indices of NaNs
        - index, a function, with signature indices= index(logical_indices),
          to convert logical indices of NaNs to 'equivalent' indices
    Example:
        >>> # linear interpolation of NaNs
        >>> nans, x= nan_helper(y)
        >>> y[nans]= np.interp(x(nans), x(~nans), y[~nans])
    """

    return np.isnan(y), lambda z: z.nonzero()[0]
# ========================================================================================================

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("-m", "--nan_corr_method", dest="nan_corr_method", default='1D')
    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( inputs.nan_corr_method )
    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()
