from astropy.io import fits
import numpy as np
import os
from IPython import embed

in_dir = './FILES/'
out_dir = './NEWFILES/'

def mod_fits_files(filename, outfilename, sourcename):
    with fits.open(filename) as hdu:
                
        print(filename)
    
        hdr = hdu[0].header
        hdr['SPECSYS'] = 'LSRK'
        hdr['VELREF'] = 257
        hdr['INSTRUME'] = 'ALMA    '
        hdr['DATAMIN'] = np.nanmin(hdu[0].data)
        hdr['DATAMAX'] = np.nanmax(hdu[0].data)
                
        hdr['OBJECT'] = sourcename
        print('Add object name:', sourcename)
        
        #flip the data axis
        newdata = np.flip(hdu[0].data, axis = 1)
        #correct the wcs to be consistent with the flip
        #Normal case with CRPIX3 at 1 (first pixel as reference)
        #This becomes the last pixel and we can just change the sign of CDELT3
        if hdr['CDELT3'] < 0:
            if hdr['CRPIX3'] == 1: #normally the case for the continuum and cubes
                hdr['CRPIX3'] = np.shape(newdata)[1]
                hdr['CDELT3'] *= -1
            else:
                hdr['CDELT3'] *= -1
                hdr['CRVAL3'] = hdr['CRVAL3'] + hdr['CDELT3'] * (hdr['CRPIX3']-1)
                hdr['CRPIX3'] = 1
        
            hdr['COMMENT'] = 'The frequency axis has been flipped using astropy to match the ALMA archive requirements (CDELT3>0) and the WCS parameters modified accordingly.'
    
        hdr['COMMENT'] = 'The reference frame has been set from TOPO to LSRK to match the archive requirement, but the negligible frequency shift (broad line and low spectral resolution) has not been applied to the WCS parameters (see also the README file).'
   
        if hdr['CDELT3'] < 0:
            print('WARNING! CDELT IS STILL NEGATIVE!')

        newhdu = fits.PrimaryHDU()
        newhdu.header = hdr
        newhdu.data = newdata  
        newhdu.writeto(outfilename, overwrite = True)

member_list = os.listdir(in_dir)

for member_dir in member_list:
    if member_dir[0] != '.':
        file_list =  os.listdir(in_dir+member_dir)
        #create the output member directory if it does no exist
        if not os.path.isdir(out_dir+member_dir):
            os.mkdir(out_dir+member_dir)
        for filename in file_list:
            sourcename = filename.split('.')[3]
            print(sourcename)
            infilename = in_dir+member_dir+'/'+filename
            outfilename = out_dir+member_dir+'/'+filename
            mod_fits_files(infilename, outfilename, sourcename)