# -*- coding: future_fstrings -*-
import stimela
import os
import unittest
import subprocess
import shutil


class TestKat7Reduce(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        unittest.TestCase.setUpClass()
        # I/O
        global INPUT
        INPUT = 'input'
        global MSDIR
        MSDIR = 'msdir'

        global MS
        MS = 'kat-7-small.ms'
        global PREFIX
        PREFIX = 'kat7_small_LBand'

        # Fields
        global GCAL
        GCAL = 'PKS2326-477'
        global TARGET
        TARGET = '1'
        global BPCAL
        BPCAL = 'PKS1934-638'

        # Reference antenna
        global REFANT
        REFANT = '0'

        # Calibration tables
        global ANTPOS_TABLE
        ANTPOS_TABLE = PREFIX + '.antpos:output'
        global BPCAL_TABLE
        BPCAL_TABLE = PREFIX + '.B0:output'
        global DELAYCAL_TABLE
        DELAYCAL_TABLE = PREFIX + '.K0:output'
        global GAINCAL_TABLE
        GAINCAL_TABLE = PREFIX + '.G0:output'
        global FLUXSCALE_TABLE
        FLUXSCALE_TABLE = PREFIX + '.fluxscale:output'

        global LABEL
        LABEL = "test_reduction"
        global OUTPUT
        OUTPUT = "output_%s" % LABEL
        global MSCONTSUB
        MSCONTSUB = MS + '.contsub'

        global SPW
        SPW = '0:100~355'
        # Calibration tables
        global LSM0
        LSM0 = PREFIX + '.lsm.html'
        global SELFCAL_TABLE1
        SELFCAL_TABLE1 = PREFIX + '.SF1:output'
        global IMAGE1
        IMAGE1 = PREFIX + 'image1:output'
        global MASK1
        MASK1 = PREFIX + 'mask1.fits'
        global IMAGE2
        IMAGE2 = PREFIX + 'image2:output'
        global nchans
        nchans = 256
        global chans
        chans = [100, 355]

        # Clean-Mask-Clean
        global imname0
        imname0 = PREFIX + 'image0'
        global maskname0
        maskname0 = PREFIX + 'mask0.fits'
        global maskname01
        maskname01 = PREFIX + 'mask01.fits'
        global imname1
        imname1 = PREFIX + 'image1'
        global corr_ms
        corr_ms = MS + '-corr.ms'
        global lsm0
        lsm0 = PREFIX + '-LSM0'
        stimela.register_globals()

    @classmethod
    def tearDownClass(cls):
        unittest.TestCase.tearDownClass()

    def tearDown(self):
        unittest.TestCase.tearDown(self)

    def setUp(self):
        unittest.TestCase.setUp(self)

    def test_end_to_end_reduction(self):
        global INPUT, OUTPUT, MSDIR, MS, LABEL
        global GAINCAL_TABLE2, FLUXSCALE_TABLE, GAINCAL_TABLE, DELAYCAL_TABLE, BPCAL_TABLE, ANTPOS_TABLE
        global REFANT, BPCAL, TARGET, GCAL, PREFIX
        global MSCONTSUB, SPW, LSM0, SELFCAL_TABLE1, corr_ms, lsm0
        global IMAGE1, IMAGE2, MASK1, nchans, chans, imname0, maskname0, maskname01, imname1

        recipe = stimela.Recipe('Test reduction script',
                                ms_dir=MSDIR, JOB_TYPE="docker", log_dir="logs")

        recipe.add('cab/casa_listobs', 'listobs', {
            "vis": MS
        },
            input=INPUT,
            output=OUTPUT,
            label='listobs:: some stats',
            time_out=300)

        recipe.add("cab/owlcat_plotelev", "plotobs", {
            "msname" : MS,
            "output-name" : "obsplot.png",
        },
            input=INPUT,
            output=OUTPUT,
            label="plotobs:: Plot elevation/azimuth vs LST/UTC")

        # It is common for the array to require a small amount of time to settle down at the start of a scan. Consequently, it has
        # become standard practice to flag the initial samples from the start
        # of each scan. This is known as 'quack' flagging
        recipe.add('cab/casa_flagdata', 'quack_flagging', {
            "vis": MS,
            "mode": 'quack',
            "quackinterval": 10.0,
            "quackmode": 'beg',
        },
            input=INPUT,
            output=OUTPUT,
            label='quack_flagging:: Quack flagging',
            time_out=300)

        # Flag the autocorrelations
        recipe.add("cab/politsiyakat_autocorr_amp", "flag_autopower", {
            "msname": MS,
            "field": "0,1,2",
            "cal_field": "0,2",
            "nrows_chunk": 5000,
            "scan_to_scan_threshold": 1.5,
            "antenna_to_group_threshold": 4,
            "nio_threads": 1,
            "nproc_threads": 32,
            },input=INPUT, output=OUTPUT, label="flag_autopower")

        recipe.add('cab/casa_flagdata', 'autocorr_flagging', {
            "vis":   MS,
            "mode":   'manual',
            "autocorr":   True,
        },
            input=INPUT,
            output=OUTPUT,
            label='autocorr_flagging:: Autocorrelations flagging',
            time_out=300)

        # Flag bad channels
        recipe.add('cab/casa_flagdata', 'badchan_flagging', {
            "vis":   MS,
            "mode":   'manual',
            "spw":   "0:113~113,0:313~313,0:369~369,0:601~607,0:204~204,0:212~212,0:594~600",

        },
            input=INPUT,
            output=OUTPUT,
            label='badchan_flagging:: Bad Channel flagging',
            time_out=300)

        recipe.add('cab/casa_clearcal', 'clearcal',
                   {
                       "vis":   MS,
                       "addmodel":   True
                   },
                   input=INPUT,
                   output=OUTPUT,
                   label='clearcal:: casa clearcal',
                   time_out=300)

        recipe.add('cab/casa_setjy', 'set_flux_scaling', {
            "vis":   MS,
            "field":   BPCAL,
            "standard":   'Perley-Butler 2010',
            "usescratch":   True,
            "scalebychan":   True,
        },
            input=INPUT,
            output=OUTPUT,
            label='set_flux_scaling:: Set flux density value for the amplitude calibrator',
            time_out=300)

        recipe.add('cab/casa_bandpass', 'bandpass_cal', {
            "vis": MS,
            "caltable": BPCAL_TABLE,
            "field": BPCAL,
            "refant": REFANT,
            "spw": SPW,
            "solint": 'inf',
            "bandtype": 'B',
            #                        "opacity"   : 0.0,
            #                        "gaincurve" : False,
        },
            input=INPUT,
            output=OUTPUT,
            label='bandpass_cal:: Bandpass calibration',
            time_out=300)

        recipe.add('cab/ragavi', 'ragavi_gains_plot_bandpass', {
                   'table': BPCAL_TABLE,
                   'gaintype': "B",
                   'htmlname': PREFIX + '_B0_amp_chan'
                   },
                   input=INPUT,
                   output=OUTPUT,
                   label='ragavi_gains_plot_bandpass:: Plot bandpass table',
                   time_out=1200
                   )

        # display the bandpass solutions. Note that in the plotcal inputs below, the amplitudes are being displayed as a function of
        # frequency channel. The parameter subplot=221 is used to display multiple plots per page (2 plots per page in the y
        # direction and 2 in the x direction). The first two commands below show the amplitude solutions (one per each polarization)
        # and the last two show the phase solutions (one per each polarization). Parameter iteration='antenna' is used to step
        # through separate plots for each antenna.
        recipe.add('cab/casa_plotcal', 'plot_bandpass_amp_R', {
            "caltable": BPCAL_TABLE,
            "poln": 'R',
            "xaxis": 'chan',
            "yaxis": 'amp',
            "field": BPCAL,
            "spw": SPW,
            "subplot": 221,
            "figfile": PREFIX + '-B0-R-amp.png',
        },
            input=INPUT,
            output=OUTPUT,
            label='plot_bandpass_amp_R:: Plot bandpass table. AMP, R',
            time_out=1200)

        # Gain calibration - amplitude and phase - first for BPCAL.
        recipe.add('cab/casa_gaincal', 'gaincal_bp', {
            "vis": MS,
            "caltable": GAINCAL_TABLE,
            "field": "{0:s},{1:s}".format(BPCAL, GCAL),
            "solint": 'inf',
            "refant": '',
            "gaintype": 'G',
            "calmode": 'ap',
            "spw": SPW,
            "solnorm": False,
            "gaintable": [BPCAL_TABLE],
            "interp": ['nearest'],
        },
            input=INPUT,
            output=OUTPUT,
            label="gaincal:: Gain calibration",
            time_out=300, 
            version=None)

        # Set fluxscale
        recipe.add('cab/casa_fluxscale', 'fluxscale', {
            "vis": MS,
            "caltable": GAINCAL_TABLE,
            "fluxtable": FLUXSCALE_TABLE,
            "reference": [BPCAL],
            "transfer": [GCAL],
            "save_result" : "fluxinfo.pickle",
            "incremental": False,
        },
            input=INPUT,
            output=OUTPUT,
            label='fluxscale:: Set fluxscale',
            time_out=300)

        # Apply calibration to BPCAL
        recipe.add('cab/casa_applycal', 'applycal_bp', {
            "vis": MS,
            "field": BPCAL,
            "gaintable": [BPCAL_TABLE, FLUXSCALE_TABLE],
            "gainfield": ['', '', BPCAL],
            "interp": ['', '', 'nearest'],
            "calwt": [False],
            "parang": False,
            "applymode": "calflag",
        },
            input=INPUT,
            output=OUTPUT,
            label='applycal_bp:: Apply calibration to Bandpass Calibrator',
            time_out=1800)

        # Flag the phase
        recipe.add("cab/politsiyakat_cal_phase", "flag_calphase", {
            "msname": MS,
            "field": ",".join(["0","1","2"]),
            "cal_field": ",".join(["0","2"]),
            "nrows_chunk": 5000,
            "data_column": "CORRECTED_DATA",
            "scan_to_scan_threshold": 1.5,
            "baseline_to_group_threshold": 4,
            "nio_threads": 1,
            "nproc_threads": 32,
            },input=INPUT, output=OUTPUT, label="flag_calphase")

        recipe.run()

        recipe = stimela.Recipe('KAT reduction script 2',
                                ms_dir=MSDIR, JOB_TYPE="docker", log_dir="logs")
        # Copy CORRECTED_DATA to DATA, so we can start uv_contsub
        recipe.add("cab/msutils", "move_corrdata_to_data", {
            "command": "copycol",
            "msname": MS,
            "fromcol": "CORRECTED_DATA",
            "tocol": "DATA",
        },
            input=INPUT, output=OUTPUT,
            label="move_corrdata_to_data::msutils",
            time_out=1800)

        os.system("rm -rf {}/{}-corr.ms".format(MSDIR, MS[:-3]))
        recipe.add('cab/casa_split', 'split_corr_data',
                   {
                       "vis":   MS,
                       "outputvis":   MS[:-3] + '-corr.ms',
                       "field":   str(BPCAL),
                       "spw":   SPW,
                       "datacolumn":   'data',
                   },
                   input=INPUT,
                   output=OUTPUT,
                   label='split_corr_data:: Split corrected data from MS',
                   time_out=1800)

        MS = MS[:-3] + '-corr.ms'

        recipe.add('cab/casa_clearcal', 'prep_split_data',
                   {
                       "vis":   MS,
                       "addmodel":   True
                   },
                   input=INPUT,
                   output=OUTPUT,
                   label='prep_split_data:: Prep split data with casa clearcal',
                   time_out=1800)

        # Clean-Mask-Clean
        imname0 = PREFIX + 'image0'
        maskname0 = PREFIX + 'mask0.fits'
        maskname01 = PREFIX + 'mask01.fits'
        imname1 = PREFIX + 'image1'

        recipe.add('cab/casa_tclean', 'image_target_field_r1', {
            "vis":   MS,
            "datacolumn":  "corrected",
            "field":   "0",
            "start":   21,  # Other channels don't have any data
            "nchan":   235 - 21,
            "width":   1,
            # Use Briggs weighting to weigh visibilities for imaging
            "weighting":   "briggs",
            "robust":   0,
            "imsize":   256,                   # Image size in pixels
            "cellsize":   "30arcsec",                      # Size of each square pixel
            "niter":   100,
            "stokes":   "I",
            "prefix":   '%s:output' % (imname1),
        },
            input=INPUT,
            output=OUTPUT,
            label="image_target_field_r1:: Image target field second round",
            time_out=300)

        recipe.add('cab/cleanmask', 'mask0', {
            "image": '%s.image.fits:output' % (imname1),
            "output": '%s:output' % (maskname0),
            "dilate": False,
            "sigma": 20,
        },
            input=INPUT,
            output=OUTPUT,
            label='mask0:: Make mask',
            time_out=1800)

        lsm0 = PREFIX + '-LSM0'
        # Source finding for initial model
        recipe.add("cab/pybdsm", "extract_init_model", {
            "image":  '%s.image.fits:output' % (imname1),
            "outfile":  '%s:output' % (lsm0),
            "thresh_pix":  25,
            "thresh_isl":  15,
            "port2tigger":  True,
        },
            input=INPUT, output=OUTPUT,
            label="extract_init_model:: Make initial model from preselfcal image",
            time_out=1800)

        # First selfcal round

        recipe.add("cab/calibrator", "calibrator_Gjones_subtract_lsm0", {
            "skymodel": "%s.lsm.html:output" % (lsm0),
            "msname": MS,
            "threads": 16,
            "column": "DATA",
            "output-data": "CORR_RES",
            "Gjones": True,
            # Ad-hoc right now, subject to change
            "Gjones-solution-intervals": [20, 0],
            "Gjones-matrix-type": "GainDiagPhase",
            "tile-size": 512,
            "field-id": 0,
        },
            input=INPUT, output=OUTPUT,
            label="calibrator_Gjones_subtract_lsm0:: Calibrate and subtract LSM0",
            time_out=1800)

        # Diversity is a good thing... lets add some DDFacet to this soup bowl
        imname = PREFIX + 'ddfacet'

        recipe.add("cab/ddfacet", "ddfacet_test",
                   {
                       "Data-MS": [MS],
                       "Output-Name": imname,
                       "Image-NPix": 256,
                       "Image-Cell": 30,
                       "Cache-Reset": True,
                       "Freq-NBand": 2,
                       "Weight-ColName": "WEIGHT",
                       "Data-ChunkHours": 10,
                       "Beam-FITSFeed": "rl",
                       "Data-Sort": True,
                       "Log-Boring": True,
                       "Deconv-MaxMajorIter": 1,
                       "Deconv-MaxMinorIter": 20,
                   },
                   input=INPUT, output=OUTPUT, shared_memory="8gb",
                   label="image_target_field_r0ddfacet:: Make a test image using ddfacet",
                   time_out=520)

        lsm1 = PREFIX + '-LSM0'
        # Source finding for initial model
        recipe.add("cab/pybdsm", "extract_init_model", {
            "image":  '%s.app.restored.fits:output' % (imname),
            "outfile":  '%s:output' % (lsm1),
            "thresh_pix":  25,
            "thresh_isl":  15,
            "port2tigger":  True,
        },
            input=INPUT, output=OUTPUT,
            label="extract_init_model:: Make initial model from preselfcal image",
            time_out=1800)

        # Stitch LSMs together
        lsm2 = PREFIX + '-LSM2'
        recipe.add("cab/tigger_convert", "stitch_lsms1", {
            "input-skymodel":   "%s.lsm.html:output" % lsm0,
            "output-skymodel":   "%s.lsm.html:output" % lsm2,
            "rename": True,
            "force": True,
            "append":   "%s.lsm.html:output" % lsm1,
        },
            input=INPUT, output=OUTPUT,
            label="stitch_lsms1::Create master lsm file",
            time_out=300)

        recipe.add("cab/calibrator", "calibrator_Gjones_subtract_lsm0", {
            "skymodel": "%s.lsm.html:output" % (lsm2),
            "msname": MS,
            "threads": 16,
            "column": "DATA",
            "output-data": "CORR_RES",
            "Gjones": True,
            # Ad-hoc right now, subject to change
            "Gjones-solution-intervals": [20, 0],
            "Gjones-matrix-type": "GainDiagPhase",
            "tile-size": 512,
            "field-id": 0,
        },
            input=INPUT, output=OUTPUT,
            label="calibrator_Gjones_subtract_lsm0:: Calibrate and subtract LSM0",
            time_out=1800)

        recipe.add('cab/casa_uvcontsub', 'uvcontsub',
                   {
                       "msname":    MS,
                       "field":    "0",
                       "fitorder":    1,
                   },
                   input=INPUT,
                   output=OUTPUT,
                   label='uvcontsub:: Subtract continuum in the UV plane',
                   time_out=1800)

        # Image HI
        recipe.add('cab/casa_clean', 'casa_dirty_cube',
                   {
                       "msname": MS + ".contsub",
                       "prefix": PREFIX,
                       "mode": 'channel',
                       "nchan": nchans,
                       "niter": 0,
                       "npix": 256,
                       "cellsize": 30,
                       "weight": 'natural',
                       "port2fits": True,
                   },
                   input=INPUT,
                   output=OUTPUT,
                   label='casa_dirty_cube:: Make a dirty cube with CASA CLEAN',
                   time_out=1800)

        recipe.add('cab/sofia', 'sofia',
                   {
                       #    USE THIS FOR THE WSCLEAN DIRTY CUBE
                       #    "import.inFile"     :   '{:s}-cube.dirty.fits:output'.format(combprefix),
                       #    USE THIS FOR THE CASA CLEAN CUBE
                       # CASA CLEAN cube
                       "import.inFile":   '{:s}.image.fits:output'.format(PREFIX),
                       "steps.doFlag": False,
                       "steps.doScaleNoise": True,
                       "steps.doSCfind": True,
                       "steps.doMerge": True,
                       "steps.doReliability": False,
                       "steps.doParameterise": False,
                       "steps.doWriteMask": True,
                       "steps.doMom0": False,
                       "steps.doMom1": False,
                       "steps.doWriteCat": True,
                       "flag.regions": [],
                       "scaleNoise.statistic": 'mad',
                       "SCfind.threshold": 4,
                       "SCfind.rmsMode": 'mad',
                       "merge.radiusX": 2,
                       "merge.radiusY": 2,
                       "merge.radiusZ": 2,
                       "merge.minSizeX": 2,
                       "merge.minSizeY": 2,
                       "port2tigger": False,
                       "merge.minSizeZ": 2,
                   },
                   input=INPUT,
                   output=OUTPUT,
                   label='sofia:: Make SoFiA mask and images',
                   time_out=1800)

        recipe.run()
