# general self-cal script for my ACA variability data

import os
import datetime
import matplotlib.pyplot as plt

# Note that 7m integrations are 1.01s, our scans are roughly 1 min in length

# general parameters
self_cal_ms='Serpens_Main_850_00_epoch1_cont_chan_avg_self-cal.ms'
basename='Serpens_Main_850_00_epoch1_self-cal'
execfile('../reduction_utils.py')

# total time ranges per epoch
#'13-Oct-2018/23:39:37.1~14-Oct-2018/00:39:46.8'
#'2019/03/06/13:04:26.9~14:05:51.7'
#2019/04/07/08:10:19.4~09:10:56.4
#2019/05/18/05:23:03.8~06:23:37.3
#2019/08/04/00:14:30.7~01:15:10.2
#2019/09/20/01:21:07.6~02:21:25.9
#2019/08/29/21:37:30.7~22:37:26.1

# Pipeline reference Antennas per epoch
# Epoch 1
# CM02, CM05, CM03, CM12, CM08, CM11, CM04, CM06, CM09, CM01
# Epoch 2
# CM02, CM12, CM08, CM06, CM09, CM11, CM04, CM01, CM07
# Epoch 3
# CM10, CM05, CM01, CM09, CM07, CM11, CM04, CM06, CM08, CM02
# Epoch 4
#CM03, CM10, CM12, CM02, CM05, CM08, CM06, CM07, CM01, CM11, CM04
# Epoch 5
# CM03, CM02, CM12, CM05, CM08, CM06, CM07, CM01, CM04
# Epoch 6
# CM03, CM10, CM12, CM05, CM08, CM07, CM01, CM11, CM04
# Epoch 7
# CM03, CM10, CM02, CM12, CM05, CM08, CM06, CM07, CM01, CM11, CM04



# calibration parameters
#refant = 'CM03, CM10, CM12' # for all epochs
refant = 'CM02, CM05, CM03, CM12, CM08, CM11, CM04, CM06, CM09, CM01'
spwmap = [0,0,0,0]

# scan timeranges for plotting, now set up automatically :)
time_ranges=[]
msmd.open(self_cal_ms)
scans = msmd.scannumbers()
for scan in scans:
    times = msmd.timesforscan(scan)
    tstart, tend = au.mjdsecToTimerange(times[0], times[-1]).split('~')
    time_ranges.append(
        (datetime.datetime.strptime(tstart, '%Y/%m/%d/%H:%M:%S.%f'),
         datetime.datetime.strptime(tend,    '%Y/%m/%d/%H:%M:%S.%f')))
msmd.done()

# Round 0 self-cal
# ------------------------------------------------------------------------------#
# save initial flags
flagmanager(vis=self_cal_ms,mode='save',versionname='before_selfcal',merge='replace')

# intitial clean with no self-cal for benchmarking, no model saving
tclean(vis=self_cal_ms,
       datacolumn='data', # NOTE COLUMN
       imagename=basename+'_no_self-cal',
       specmode='mfs',
       gridder='standard',
       imsize=[256, 256],
       cell='0.3arcsec',
       weighting='briggs',
       robust=0.5,
       pblimit=0.01,
       niter=9999,
       interactive=True)

# masks for measuring SNR
disk_mask='circle[[128pix, 128pix], 12arcsec]'
noise_mask='annulus[[128pix, 128pix], [16arcsec, 32arcsec]]'

# measure SNR with no self-cal
estimate_SNR(imagename=basename+'_no_self-cal.image',
             disk_mask=disk_mask,
             noise_mask=noise_mask)
#Serpens_Main_850_00_epoch1_self-cal_no_self-cal.image
#Beam 5.854 arcsec x 2.760 arcsec (-82.48 deg)
#Flux inside disk mask: 5828.85 mJy
#Peak intensity of source: 3579.30 mJy/beam
#rms: 3.73e+01 mJy/beam
#Peak SNR: 95.84



# initial clean for first round of self-cal
tclean(vis=self_cal_ms,
       datacolumn='data', #NOTE COLUMN!!! ORIGINAL DATA!!!
       imagename=basename+'_p0',
       specmode='mfs',
       gridder='standard',
       imsize=[256, 256],
       cell='0.3arcsec',
       weighting='briggs',
       robust=0.5,
       pblimit=0.01,
       niter=9999, #  11 iter used
       savemodel='modelcolumn',
       interactive=True)

# Remaining rounds of self-cal
# ----------------------------------------------------------#

# self-cal round 1 params
tablename=basename+'_pcal1'
gaintype='G' # per pol solns
calmode='p' # phase only
combine='' # solns per spw and scan
solint='inf' # full scans
minblperant=5
figfile=basename+'_p1_scan_{}_solns.png'
versionname='after_pcal1'
imagename=basename+'_p1'
# iter cleaned: 25
# stats:
#Serpens_Main_850_00_epoch1_self-cal_p1.image
#Beam 5.778 arcsec x 2.759 arcsec (-82.71 deg)
#Flux inside disk mask: 3613.28 mJy
#Peak intensity of source: 3387.86 mJy/beam
#rms: 7.79e+01 mJy/beam
#Peak SNR: 43.51





# self-cal round 2 params
tablename=basename+'_pcal2'
gaintype='T' # combine pol solns
calmode='p' # phase only
combine='' # solns per spw and scan
solint='20.2s' # roughly a third of a scan, 20 ints
minblperant=5
figfile=basename+'_p2_scan_{}_solns.png'
versionname='after_pcal2'
imagename=basename+'_p2'
# iter cleaned: 173
# stats:
#Serpens_Main_850_00_epoch1_self-cal_p2.image
#Beam 5.778 arcsec x 2.759 arcsec (-82.71 deg)
#Flux inside disk mask: 5566.29 mJy
#Peak intensity of source: 3582.33 mJy/beam
#rms: 2.11e+01 mJy/beam
#Peak SNR: 170.17



# self-cal round 3 params
tablename=basename+'_pcal3'
gaintype='T' # combine pol solns
calmode='p' # phase only
combine='' # solns per spw and scan
solint='5.05s' # roughly a third of a scan, 5 ints
minblperant=5
figfile=basename+'_p3_scan_{}_solns.png'
versionname='after_pcal3'
imagename=basename+'_p3'
# iter cleaned: 572
# stats:
#Serpens_Main_850_00_epoch1_self-cal_p3.image
#Beam 5.774 arcsec x 2.759 arcsec (-82.72 deg)
#Flux inside disk mask: 6356.85 mJy
#Peak intensity of source: 3661.44 mJy/beam
#rms: 1.03e+01 mJy/beam
#Peak SNR: 355.05



#
# Copy and paste the below *after* setting per-round parameters above
#

# solve for solutions
gaincal(vis=self_cal_ms,
        caltable=tablename,
        gaintype=gaintype,
        refant=refant,
        calmode=calmode,
        combine=combine,
        solint=solint, 
        minsnr=3.0,
        minblperant=minblperant)

# plot solutions for each antenna, all scans
plotcal(caltable=tablename,
        xaxis='time',
        yaxis='phase',
        iteration='antenna',
        subplot=431,
        plotrange=[0,0,-60,60],
        )

# plot solutions for each antenna and  scan
for i, trange in enumerate(time_ranges):                                                                                                                                                                
    plotcal(caltable=tablename,
            xaxis='time',
            yaxis='phase',
            iteration='antenna',
            subplot=431,
            plotrange=[0,0,-60,60],
            )
    # stupid workaround for setting subplot timeranges
    fig=plt.gcf()
    axes=fig.get_axes()
    for ax in axes:
        ax.set_xlim(trange)
    fig.savefig(figfile.format(i+1))
    raw_input("Press Enter to continue...")

        
# apply solutions
applycal(vis=self_cal_ms,
         spwmap=spwmap,
         gaintable=[tablename],
         gainfield='',
         calwt=False,
         flagbackup=False,
         interp='linear')

# save flags
flagmanager(vis=self_cal_ms,mode='save',versionname=versionname)

# clean deeper
tclean(vis=self_cal_ms,
       datacolumn='corrected', # NOTE CHANGE OF COLUMN!!! MUST CLEAN CORRECTED!!!
       imagename=imagename,
       specmode='mfs',
       gridder='standard',
       imsize=[256, 256],
       cell='0.3arcsec',
       weighting='briggs',
       robust=0.5,
       pblimit=0.01,
       niter=9999,
       savemodel='modelcolumn',
       interactive=True)

# measure SNR
estimate_SNR(imagename=imagename+'.image',
             disk_mask=disk_mask,
             noise_mask=noise_mask)
#Serpens_Main_850_00_epoch1_self-cal_p1.image
#Beam 5.778 arcsec x 2.759 arcsec (-82.71 deg)
#Flux inside disk mask: 5955.05 mJy
#Peak intensity of source: 3607.80 mJy/beam
#rms: 1.52e+01 mJy/beam
#Peak SNR: 236.92


# ---------------------------------------#

# reset to beginning of round if needed:
# delete clean images
for ext in ['.image','.mask','.model','.image.pbcor','.psf','.residual','.pb','.sumwt']:
    rmtables(imagename + ext)
# delete calibration table
rmtables(tablename)
# delete plots
for i, trange in enumerate(time_ranges):                                                                                                                                                                
    os.system('rm ' + figfile.format(i+1))
# restore flags to last round, ALTER AS NEEDED!
flagmanager(vis=self_cal_ms, mode='restore',versionname='before_selfcal')
# copy data column to corrected column, set model column to 1s
clearcal(vis=self_cal_ms)
# clear model column
delmod(vis=self_cal_ms,otf=True)
# *REAPPLY last rounds calibration* (not needed if you are still on round 1 obviously)
applycal(vis=self_cal_ms,
         spwmap=spwmap,
         gaintable=[], # LAST ROUND CAL TABLE HERE!!!!
         gainfield='',
         calwt=False,
         flagbackup=False,
         interp='linear')

