import numpy
import os
import sys
from pyrap.tables import table
import subprocess

sys.path.append("/scratch/stimela")

utils = __import__('utils')

CONFIG = os.environ["CONFIG"]
INPUT = os.environ["INPUT"]
OUTPUT = os.environ["OUTPUT"]
MSDIR = os.environ["MSDIR"]
CODE = "/scratch/code"

cab = utils.readJson(CONFIG)
binary = cab['binary']
parameters = cab['parameters']

jdict = {}
for param in parameters:
    name = param['name']
    value = param['value']
    if value is None:
        continue
    if value is False:
        value = 0
    elif value is True:
        value = 1

    jdict[name] = value

msname = jdict['msname']
THREADS = jdict.pop('threads', 2)
msbase = os.path.basename(msname)[:-3]
prefix = jdict.pop('prefix', None) or '{0}/{1}'.format(OUTPUT, msbase)
params = {}

# options for writing flags
writeflags = jdict.pop("write-flags-to-ms", None)
if writeflags:
    params["ms_sel.ms_write_flags"] = 1
    params["ms_sel.ms_fill_legacy_flags"] = 1 if jdict.pop(
        "fill-legacy-flags", False) else 0

write_flagset = jdict.pop("write-flagset", None)
if write_flagset:
    params["ms_wfl.write_bitflag"] = write_flagset
    params["ms_sel.ms_write_flag_policy"] = "'add to set'" if jdict.pop(
        "write-flag-policy", False) else "'replace set'"

# Read flags options
readflagsets = jdict.pop("read-flagsets", False)
if readflagsets:
    params["ms_rfl.read_flagsets"] = readflagsets

params['ms_sel.ms_read_flags'] = 1 if jdict.pop(
    "read-flags-from-ms", False) else 0
params["ms_rfl.read_legacy_flags"] = 1 if jdict.pop(
    "read-legacy-flags", False) else 0

params["ms_sel.msname"] = msname
field_id = jdict.pop("field-id", 0)
spw_id = jdict.pop("spw-id", 0)
params["ms_sel.tile_size"] = jdict.pop("tile-size", 16)
params["ms_sel.ddid_index"] = spw_id
params["ms_sel.field_index"] = field_id

TDL = jdict.pop("tdlconf", None) or CODE + "/tdlconf.profiles"
SECTION = jdict.pop("section", None) or "stefcal"
skymodel = jdict.pop("skymodel", None)
beam_files_pattern = jdict.pop("beam-files-pattern", False)
jones_type = jdict.pop("jones-implementation", "Gain2x2")

column = jdict.pop("column", "DATA")
outcol = jdict.pop("output-column", "CORRECTED_DATA")

params["ms_sel.input_column"] = column
params["ms_sel.output_column"] = outcol
params["tiggerlsm.filename"] = skymodel
params["tiggerlsm.lsm_subset"] = jdict.get("subset", "all")
params["do_output"] = jdict.pop("output-data", "CORR_RES")
saveconf = jdict.pop('save-config', None)
params['ms_sel.ms_corr_sel'] = "'{}'".format(jdict.pop('correlations', '2x2'))

label = jdict.pop("label", None)

model_column = jdict.pop("model-column", 'MODEL_DATA')

gjones = jdict.pop("Gjones", False)
if gjones:

    time_smooth, freq_smooth = jdict.get("Gjones-smoothing-intervals", (1, 1))
    time_int, freq_int = jdict.get("Gjones-solution-intervals", (1, 1))
    mode = 'apply' if jdict.get('Gjones-apply-only', False) else 'solve-save'

    gjones_gains = jdict.pop('Gjones-gain-table', None) or "{0}/{1}{2}.gain.cp".format(
        OUTPUT, msbase, "-%s" % label if label else "")
    params.update({
        "stefcal_gain.enabled": 1,
        "stefcal_gain.mode": mode,
        "stefcal_gain.reset": 0 if mode == "apply" else 1,
        "stefcal_gain.implementation": jones_type,
        "stefcal_gain.timeint": time_int,
        "stefcal_gain.freqint": freq_int,
        "stefcal_gain.flag_ampl":   jdict.get("Gjones-ampl-clipping", 0),
        "stefcal_gain.flag_chisq":   jdict.get("Gjones-chisq-clipping", 0),
        "stefcal_gain.flag_chisq_threshold":   jdict.get("Gjones-thresh-sigma", 10),
        "stefcal_gain.flag_ampl_low":   jdict.get("Gjones-ampl-clipping-low", 0.3),
        "stefcal_gain.flag_ampl_high":   jdict.get("Gjones-ampl-clipping-high", 2),
        "stefcal_gain.implementation":   jdict.get("Gjones-matrix-type", "Gain2x2"),
        "stefcal_gain.table": gjones_gains,
    })

bjones = jdict.pop("Bjones", False)
if bjones:

    time_smooth, freq_smooth = jdict.get("Bjones-smoothing-intervals", (1, 0))
    time_int, freq_int = jdict.get("Bjones-solution-intervals", (1, 0))
    mode = 'apply' if jdict.get('Bjones-apply-only', False) else 'solve-save'

    bjones_gains = jdict.pop('Bjones-gain-table', None) or "{0}/{1}{2}.gain1.cp".format(
        OUTPUT, msbase, "-%s" % label if label else "")
    params.update({
        "stefcal_gain1.enabled": 1,
        "stefcal_gain1.label": 'B',
        "stefcal_gain1.mode": mode,
        "stefcal_gain1.reset": 0 if mode == "apply" else 1,
        "stefcal_gain1.implementation": jones_type,
        "stefcal_gain1.timeint": time_int,
        "stefcal_gain1.freqint": freq_int,
        "stefcal_gain1.flag_ampl":   jdict.get("Bjones-ampl-clipping", 0),
        "stefcal_gain1.flag_chisq":   jdict.get("Bjones-chisq-clipping", 0),
        "stefcal_gain1.flag_chisq_threshold":   jdict.get("Bjones-thresh-sigma", 10),
        "stefcal_gain1.flag_ampl_low":   jdict.get("Bjones-ampl-clipping-low", 0.3),
        "stefcal_gain1.flag_ampl_high":   jdict.get("Bjones-ampl-clipping-high", 2),
        "stefcal_gain1.implementation":   jdict.get("Bjones-matrix-type", "Gain2x2"),
        "stefcal_gain1.table": bjones_gains,
    })

beam = jdict.pop("Ejones", False)
if beam and beam_files_pattern:
    params.update({
        "me.e_enable": 1,
        "me.p_enable": 1,
        "me.e_module": "Siamese_OMS_pybeams_fits",
        "me.e_all_stations": 1,
        "pybeams_fits.sky_rotation":  1 if jdict.pop('parallactic-angle-rotation', False) else 0,
        "pybeams_fits.l_axis": jdict.pop("beam-l-axis", "L"),
        "pybeams_fits.m_axis": jdict.pop("beam-m-axis", "M"),
        "pybeams_fits.filename_pattern": "'{}'".format(beam_files_pattern),
    })


ddjones = jdict.pop("DDjones", False)
if ddjones:
    time_int, freq_int = jdict.pop("DDjones-solution-intervals", (1, 1))
    time_smooth, freq_smooth = jdict.pop("DDjones-smoothing-intervals", (1, 1))

    mode = 'apply' if jdict.get('DDjones-apply-only', False) else 'solve-save'

    ddjones_gains = jdict.pop('DDjones-gain-table', None) or "{0}/{1}{2}.diffgain.cp".format(
        OUTPUT, msbase, "-%s" % label if label else "")
    params.update({
        "stefcal_diffgain.enabled": 1,
        "stefcal_diffgain.reset": 0 if mode == "apply" else 1,
        "stefcal_diffgain.flag_ampl": 0,
        "stefcal_diffgain.flag_chisq": 1,
        "stefcal_diffgain.flag_chisq_threshold": 5,
        "stefcal_diffgain.freqint": freq_int,
        "stefcal_diffgain.freqsmooth": freq_smooth,
        "stefcal_diffgain.implementation": jones_type,
        "stefcal_diffgain.label": jdict.pop("DDjones-tag", "dE"),
        "stefcal_diffgain.max_diverge": 1,
        "stefcal_diffgain.mode": mode,
        "stefcal_diffgain.niter": 50,
        "stefcal_diffgain.omega": 0.5,
        "stefcal_diffgain.quota": 0.95,
        "stefcal_diffgain.table": ddjones_gains,
        "stefcal_diffgain.timeint": time_int,
        "stefcal_diffgain.timesmooth": time_smooth,
        "stefcal_diffgain.implementation":   jdict.get("DDjones-matrix-type", "Gain2x2"),
    })

ifrjones = jdict.pop("DDjones", False)
if ifrjones:
    ifrjones_gains = jdict.pop('IFRjones-gain-table', None) or "{0}/{1}{2}.ifrgain.cp".format(
        OUTPUT, msbase, "-%s" % label if label else "")
    mode = 'apply' if jdict.get('IFRjones-apply-only', False) else 'solve-save'

    params.update({
        "stefcal_ifr_gain_mode": mode,
        "stefcal_ifr_gains": 1,
        "stefcal_ifr_gain_reset": 0 if mode == "apply" else 1,
        "stefcal_reset_ifr_gains": 0 if mode == "apply" else 1,
        "stefcal_ifr_gain.table": ifrjones_gains,
    })

makeplots = jdict.pop("make-plots", False)

gjones_plotprefix = prefix+"-gjones_plots"
bjones_plotprefix = prefix+"-bjones_plots"
ddjones_plotprefix = prefix+"-ddjones_plots"
ifrjones_plotprefix = prefix+"-ifrjones_plots"


def run_meqtrees(msname):

    prefix = ["--mt %d -c %s [%s]" % (THREADS, TDL, SECTION)]
    suffix = ["%s/Calico/calico-stefcal.py =stefcal" %
              os.environ["MEQTREES_CATTERY_PATH"]]
    options = {}
    options.update(params)
    if jdict.pop("add-vis-model", 0):
        options["read_ms_model"] = 1
        options["ms_sel.model_column"] = model_column

    taql = jdict.get('data-selection', None)
    if taql:
        options["ms_sel.ms_taql_str"] = taql

    args = []
    for key, value in options.iteritems():
        if isinstance(value, str) and value.find(' ') > 0:
            value = '"{:s}"'.format(value)
        args.append('{0}={1}'.format(key, value))

    args = prefix + args + suffix

    utils.xrun(cab['binary'], args +
               ['-s {}'.format(saveconf) if saveconf else ''])

    print("MeqTrees Done!")
    # now plot the gains
    if makeplots:
        print("Preparing to make gain plots")
        import Owlcat.Gainplots as plotgains
        feed_tab = table(msname+"/FEED")
        print("Extracting feed type from MS")
        feed_type = set(feed_tab.getcol("POLARIZATION_TYPE")['array'])
        feed_type = "".join(feed_type)
        print("Feed type is [%s]" % feed_type)

        if feed_type.upper() in ["XY", "YX"]:
            feed_type = "XY"
        else:
            feed_type = "RL"

        if gjones:
            print("Making Gain plots (G)...")
            plotgains.make_gain_plots(
                gjones_gains, prefix=gjones_plotprefix, feed_type=feed_type)

        if bjones:
            print("Making Gain plots (B)...")
            plotgains.make_gain_plots(
                bjones_gains, gain_label='B', prefix=bjones_plotprefix, feed_type=feed_type)

        if ddjones:
            print("Making differential gain plots...")
            plotgains.make_diffgain_plots(
                ddjones_gains, prefix=ddjones_plotprefix, feed_type=feed_type)

        if ifrjones:
            print("Making IFR gain plots...")
            plotgains.make_ifrgain_plots(
                ifrjones_gains, prefix=ifrjones_plotprefix, feed_type=feed_type)

    sys.exit(0)


run_meqtrees(msname)
