import os

# Name of calibrated dataset
visfile     = 'uid___A002_Xf396d6_X45bb.ms'

"""
CONTINUUM IMAGING
"""
# Identified continuum (line-free) frequency ranges
contchans   = '17:86.9876873757~88.6908123757GHz,21:99.1714062009~100.9370312009GHz,'\
              '23:100.7279366207~102.4935616207GHz,25:88.7249991207~89.4632803707GHz;'\
              '89.8656241207~90.5374991207GHz'

"""
Image the continuum emission for the science target.
First we will make a so-called 'dirty image' (0 clean iterations).
This will help us to get a first look at the data and refine our parameters for
the full imaging.
"""
if not os.path.exists('PJ113921.7.cont.dirty.image'):
    tclean( vis            = visfile,
            imagename      = 'PJ113921.7.cont.dirty',
            field          = 'PJ113921.7',
            spw            = contchans,
            specmode       = 'cont',
            imsize         = [1250, 1250],
            cell           = '0.09arcsec',
            deconvolver    = 'hogbom',
            niter          = 0,
            weighting      = 'briggs',
            robust         = 0.5,
            gridder        = 'standard',
            pbcor          = False,
            interactive    = False
            )

"""
Now we will go ahead and clean the data (> 0 iterations).
Note that the no. iterations has increased from 0 -> 1000000, a cleaning
threshold has been specified, and automasking parameters have been added.

Having inspected the dirty image created previously, we can see that the
emission is concentrated at the centre of the field. We therefore restrict the
size of our cleaned image to 150x150 pixels for efficiency. Feel free to play
with this and image the full field.
"""
if not os.path.exists('PJ113921.7.cont.image'):
    tclean( vis                 = visfile,
            imagename           = 'PJ113921.7.cont',
            field               = 'PJ113921.7',
            phasecenter         = 'ICRS 11:39:21.728 20.24.51.838',
            spw                 = contchans,
            specmode            = 'cont',
            imsize              = [150, 150],
            cell                = '0.09arcsec',
            deconvolver         = 'hogbom',
            niter               = 1000000,
            weighting           = 'briggs',
            robust              = 0.5,
            gridder             = 'standard',
            pbcor               = True,
            threshold           = '2.5e-05Jy',
            usemask             = 'auto-multithresh',
            growiterations      = 75,
            sidelobethreshold   = 2.0,
            noisethreshold      = 4.0,
            lownoisethreshold   = 2.5,
            minbeamfrac         = 0.1,
            interactive         = False
            )

### _____________________________________ ###

"""
LINE IMAGING
"""

"""
Split out the spectral window containing the CO (3-2) emission (SPW 25).
This is the SPW that we will image. It's not necessary to do this step, but it
reduces the file-size significantly by discarding the irrelevant data.
"""
if not os.path.exists(visfile+'.split'):
    split(  vis             = visfile,
            outputvis       = visfile+'.split',
            field           = 'PJ113921.7',
            spw             = '25',
            datacolumn      = 'corrected'
            )

"""
Subtract continuum emission level to prepare for line imaging.
When we image the CO line data, we only want the molecular line emission.
The data contain both this discrete line emission and the the dust
continuum emission. This dust continuum emission needs to be removed across the
frequency range in order to image the CO-only emission.

First we specify the continuum frequency ranges in the CO spectral window.
Note that the SPW has been reindexed to 0 (instead of 25) due to splitting.
"""
contchans_CO    = '0:88.7249991207~89.4632803707GHz;89.8656241207~90.5374991207GHz'

if not os.path.exists(visfile+'.split.contsub'):
    uvcontsub(  vis             = visfile+'.split',
                field           = 'PJ113921.7',
                spw             = '0',
                fitspw          = contchans_CO,
                fitorder        = 1,
                excludechans    = False,
                want_cont       = False
                )

"""
As we did for the continuum, we will first make a 'dirty image' of the CO.
"""
if not os.path.exists('PJ113921.7.CO.cube.dirty.image'):
    tclean(  vis            = visfile+'.split.contsub',
             imagename      = 'PJ113921.7.CO.cube.dirty',
             field          = 'PJ113921.7',
             spw            = '0',
             restfreq       = '345.79599GHz',
             specmode       = 'cube',
             imsize         = [1250, 1250],
             cell           = '0.09arcsec',
             deconvolver    = 'hogbom',
             niter          = 0,
             weighting      = 'briggsbwtaper',
             robust         = 0.5,
             gridder        = 'standard',
             pbcor          = True,
             interactive    = False,
             )

"""
Now we will clean the CO emission.
Again, we will refine the parameters based on the dirty image.
The image size is again restricted to 150x150 pixels.
Since this is a 3D cube (unlike the 2D continuum), we can also restrict the
number of frequency channels that we image. Based on the dirty cube, we see that
(roughly) channels 100 - 200 contain the bulk of the emission, and so we only
image those channels in order to speed up the process.
"""
if not os.path.exists('PJ113921.7.CO.cube.image'):
    tclean(  vis                = visfile+'.split.contsub',
             imagename          = 'PJ113921.7.CO.cube',
             field              = 'PJ113921.7',
             phasecenter        = 'ICRS 11:39:21.728 20.24.51.838',
             spw                = '0',
             restfreq           = '345.79599GHz',
             specmode           = 'cube',
             imsize             = [150, 150],
             cell               = '0.09arcsec',
             deconvolver        = 'hogbom',
             niter              = 1000000,
             weighting          = 'briggsbwtaper',
             robust             = 0.5,
             gridder            = 'standard',
             pbcor              = True,
             threshold          = '0.001Jy',
             nchan              = 100,
             start              = 200,
             width              = 1,
             interactive        = False,
             usemask            = 'auto-multithresh',
             sidelobethreshold  = 2.0,
             noisethreshold     = 4.0,
             lownoisethreshold  = 2.5,
             minbeamfrac        = 0.1,
             growiterations     = 75
             )

"""
Final step: export all cleaned images to FITS format
"""
images = ['PJ113921.7.cont.image',
          'PJ113921.7.CO.cube.image']

for im in images:
    if not os.path.exists(im+'.pbcor.fits'):
        exportfits(imagename = im+'.pbcor',
                   fitsimage = im+'.pbcor.fits',
                   dropdeg   = True
                   )

        exportfits(imagename = im,
                   fitsimage = im+'.fits',
                   dropdeg   = True
                   )