import numpy as np
import os

#====== FUNCTION DEFINITION ====#

def getTargetSciSPW(thisMS):
    """ Function (to be run in CASA) to get the science spectral windows (so ignoring #CH_AVG SPWs) for ALMA data. Returns science spectral window IDs as a comma separated string ready for use in e.g. split.
    Copy this function definition into your CASA script and run as mySPWstring =  getTargetSciSPW(myMS)"""

    msmd.open(thisMS)
    targSPWs = msmd.spwsforintent('OBSERVE_TARGET#ON_SOURCE')
    chavgSPWs = msmd.almaspws(chavg=True,wvr=True)
    msmd.close()
    msmd.done()

    sciSPWstring = ','.join(map(str,np.setdiff1d(targSPWs,chavgSPWs)))

    return sciSPWstring

#===== END FUNCTION DEFINITION ====#


#===== MAIN CODE ====#
#--- Initial parameters ---#
myMS = ''      #--- This is a final product produced by re-running script for PI
target = ''    #--- This is the target name as listed in the data.

#--- Parameters for tutees to fill in ---#
contSPW = ''  #--- channels with only continnum emission. A string of format 'spw:X~Y' 
myCell = ''   #--- cell/pixel size of your images. A string of format 'X.xxxarcsec' 
myImsize =    #--- number of x by y pixels in your image. Format integer 

myThreshold = '' #--- 3x your calculated theoretical threshold. A string of format 'Y.ymJy'

startCO =                #--- Start Channel of CO cube. Format integer
nChanCO =                #--- Number of channels to image in CO cube. Format integer 
myThreshCO = ''          #--- CO channel noise threshold. A string of format 'Y.ymJy'

startCS =                #--- Start Channel of 13CS cube . Format integer
nChanCS =                #--- Number of channels to image in 13CS cube 
myThreshCS = ''          #--- 13CS channel noise threshold. A string of format 'Y.ymJy

myLineChans = '' #Line emission channel range. A string of format 'X~Y'

#--- Logic switches ---#
""" Set the step you wish to run to = True and all others to = False """
step0 = False        #--- Run a listobs on the measurementSet myMS
step1 = False        #--- Find and split out science Spectral Windows (SPWs)
step2 = False        #--- Inspect the data in PlotsMS
step3 = False        #--- Image the continuum.
step4 = False        #--- Image the CO and 13^CS line
step5 = False        #--- Moment maps



#-------------------#
#---- THE STEPS ----#

#-------------------------------------------------#
#--- STEP 0: Make listobs of the calibrated MS ---#
if step0:
    print("\n!!!STEP 0: Make listobs of the calibrated measurement set (MS).!!!")
    print("This step makes a listobs file of our calibrated MS. The listobs file contains useful information about the observation setup e.g. frequency settings, time on source, number of antennas and more. Making a listob at this stages means we can see what data in this MS we wish to retain.")
    
    if os.path.exists(myMS):
        print(f">> Generating listobs file of {myMS}")
        listobs(myMS,listfile=myMS+'.listobs')

        print(">> You can now inspect the listobs file in a standard text editor.")
    else:
        print(f">> You do not seem to have {myMS} in the current directory, skipping the rest of this step.")

#---------------------------------------------------------#
#--- STEP 1: Split out science target and science SPWs ---#
if step1:
    print("\n!!! STEP 1: Split out the science target and the science Spectral Windows (SPWs).!!!")
    print(f"{myMS} contains lots of data we won't actually need during imaging the science target. For example scans on the observed calibrators and SPWs which are used only for calibration. To give us a smaller MS to work with we will first split out from {myMS}, the target source ({target}) and the science SPWs.")

    scispws = getTargetSciSPW(myMS)
    print(f">> scispws: {scispws}")

    split(vis=myMS,
        intent = 'OBSERVE_TARGET#ON_SOURCE',
        datacolumn='corrected',
        field=target,
        spw = scispws,
        outputvis=target+'_TM2.split.cal')

    print(f">> New MS created: {target}_TM2.split.cal")
    
    listobs(target+'_TM2.split.cal',listfile=target+'_TM2.split.cal.listobs')
    
    print(f">> New listobs file generated: {target}_TM2.split.cal.listobs")


#-------------------------------------------------------------------#
#--- STEP 2: Inspect the data in plotMS & Subtract the continuum ---#
if step2:
    print("\n!!! STEP 2: Inspect the data in plotMS !!!\n")
    print("This is a hands-on step, please see the instructions on the tutorial webpage. Once you have inspected the data please fill in the contSPW parameter in the 'Initial parameters' section of this script.")
    
    print(">> Is contSPW set? [y/n]")
    answerS2 = input()
    if answerS2 == 'y':
        doContSub = True
    else:
    	print("  >> Please, fill in this parameter before proceeding.")
    	doContSub = False
    
    if doContSub:    
    	print("\n!!!STEP 2: Subtract the continuum!!!\n")        
    	print("We need to generate a continuum subtracted MS. For this we give CASA a range of free-line channels (contSPW) to fit the continuum to within the task uvcontsub. We should use this continuum subtracted MS for the line imaging later.")

    	print(">> Generating continuum subtracted MS.")
    	uvcontsub(vis = target+'_TM2.split.cal',
        fitspw = contSPW)
    	print(f">> Continuum subtracted MS created: {target}_TM2.split.cal.contsub.")


#-------------------------------#
#--- STEP 3: Image continuum ---#
if step3:
    print("\n!!!STEP 3: Image continuum!!!\n")
    print(f"Here we will image the spectral line free continuum emission of the target source {target}. Please read STEP 3 of the tutorial, calculate and input the required parameters into the 'Initial parameters' section of this script.")
    
    print(">> Are contSPW, myCell, myImsize and myThreshold set? [y/n]")
    answerS3 = input()
 
    if answerS3 == 'y':
        doContClean = True
    
    else:
        print("  >> Please, fill in these parameters before proceeding.")
        doContClean = False

    if doContClean:
        print("  >> Cleaning Z-CMa continuum.")
        #uncomment line below if you want to delete existing files with the same base name 
    	#os.system('rm -rf  '+target+'_all_spw_continuum.*')
    
        tclean(vis = target+'_TM2.split.cal',
          imagename = target+'_all_spw_continuum',
          field = target,
          spw=contSPW,
          threshold=myThreshold,
          cell=[myCell],
          imsize=[myImsize,myImsize], #should cover about the FoV 
          niter=10000,
          deconvolver='multiscale',
          scales=[0,5,15],# inpixels equivalent to point source and 1, 3 times the beam.
          weighting = 'briggs',
          robust=0.5,
          pbcor=True, 
          specmode='mfs',     
          interactive=True
          )

#---------------------------#
#--- STEP 4: Image lines ---#
if step4:
    print("\n!!!STEP 4: Image lines!!!\n")
    print("Now we image the CO line, we will use CASA's automasking capability.")

    print("  >> Cleaning CO emission in Z-CMa.")

    #uncomment line below if you want to delete existing files with the same base name 
    #os.system('rm -rf  '+target+'_CO_cube.*')
    
    tclean(vis = target+'_TM2.split.cal.contsub',
        imagename = target+'_CO_cube',
        field = target,
        stokes = 'I',
        spw = '8', 
        outframe = 'LSRK',
        restfreq = '230.538GHz',
        specmode = 'cube',
        width = 1,
        nchan = nChanCO,
        start = startCO, 
        cell=[myCell],
        imsize=[myImsize,myImsize], #covers the FoV 
        deconvolver = 'multiscale',
        scales=[0,5,15], # in pixels, equivalent to point source and 1, 3 times the beam.
        niter = 10000000,
        weighting = 'briggs',
        robust = 0.5,
        usemask = 'auto-multithresh', 
        pbcor = True,
        threshold = myThreshCO, 
        interactive = True
        )

    print("  >> Now Cleaning 13CS emission in Z-CMa.")
    #uncomment line below if you want to delete existing files with the same base name 
    #os.system('rm -rf  '+target+'_13CS_cube.*')
    
    tclean(vis = target+'_TM2.split.cal.contsub',
        imagename = target+'_13CS_cube',
        field = target,
        stokes = 'I',
        spw = '9', 
        outframe = 'LSRK',
        restfreq = '231.2206GHz',
        specmode = 'cube',
        width = 1,
        nchan = nChanCS,
        start = startCS, 
        cell=[myCell],
        imsize=[myImsize,myImsize], #covers the FoV 
        deconvolver = 'multiscale',
        scales=[0,5,15],# inpixels equivalent to point source and 1, 3 times the beam.
        niter = 10000000,
        weighting = 'briggs',
        robust = 0.5,
        usemask = 'auto-multithresh',
        pbcor = True,
        threshold = myThreshCS, 
        interactive = True
        )

#-------------------------------#
#--- STEP 5: Moment Analysis ---#
if step5:
    print('\n!!! STEP 5: Moment Analysis !!!\n')
    print(f"Finally we generate 0th, 1st and 2nd order image moments of the CO emission from {target}. This gives us integrated line intensity, the velocity fields and the velocity dispersion of the CO.")

    print(f"  >> Now Generating CO moments 0,1 & 2 for {target}.")
    
    #uncomment line below if you want to delete existing files with the same base name 
    #os.system('rm -rf  '+target+'_CO_mom*')
    immoments(imagename  = target+'_CO_cube.image',
        moments =[0,1,2],
        axis='spectral',
        chans = myLineChans,
        excludepix = [-100000.0,0.2],
        outfile = target+'_CO_mom'
        )
        
#--- END THE STEPS ---#
#---------------------#