#!/usr/bin/env python
#=============================================================================#
#                                                                             #
# NAME:     do_RMsynth_1D.py                                                  #
#                                                                             #
# PURPOSE: API for runnning RM-synthesis on an ASCII Stokes I, Q & U spectrum.#
#                                                                             #
# MODIFIED: 16-Nov-2018 by J. West                                            #
# MODIFIED: 23-October-2019 by A. Thomson                                     #
#                                                                             #
#=============================================================================#
#                                                                             #
# The MIT License (MIT)                                                       #
#                                                                             #
# Copyright (c) 2015 - 2018 Cormac R. Purcell                                 #
#                                                                             #
# Permission is hereby granted, free of charge, to any person obtaining a     #
# copy of this software and associated documentation files (the "Software"),  #
# to deal in the Software without restriction, including without limitation   #
# the rights to use, copy, modify, merge, publish, distribute, sublicense,    #
# and/or sell copies of the Software, and to permit persons to whom the       #
# Software is furnished to do so, subject to the following conditions:        #
#                                                                             #
# The above copyright notice and this permission notice shall be included in  #
# all copies or substantial portions of the Software.                         #
#                                                                             #
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR  #
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,    #
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE #
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER      #
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING     #
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER         #
# DEALINGS IN THE SOFTWARE.                                                   #
#                                                                             #
#=============================================================================#

import sys
import os
import time
import traceback
import json
import math as m
import numpy as np
import matplotlib.pyplot as plt

from RMutils.util_RM import do_rmsynth
from RMutils.util_RM import do_rmsynth_planes
from RMutils.util_RM import get_rmsf_planes
from RMutils.util_RM import measure_FDF_parms
from RMutils.util_RM import measure_qu_complexity
from RMutils.util_RM import measure_fdf_complexity
from RMutils.util_misc import nanmedian
from RMutils.util_misc import toscalar
from RMutils.util_misc import create_frac_spectra
from RMutils.util_misc import poly5
from RMutils.util_misc import MAD
from RMutils.util_plotTk import plot_Ipqu_spectra_fig
from RMutils.util_plotTk import plot_rmsf_fdf_fig
from RMutils.util_plotTk import plot_complexity_fig
from RMutils.util_plotTk import CustomNavbar
from RMutils.util_plotTk import plot_rmsIQU_vs_nu_ax

if sys.version_info.major == 2:
    print('RM-tools will no longer run with Python 2! Please use Python 3.')
    exit()

C = 2.997924538e8 # Speed of light [m/s]

#-----------------------------------------------------------------------------#
def run_rmsynth(data, polyOrd=3, phiMax_radm2=None, dPhi_radm2=None,
                nSamples=10.0, weightType="variance", fitRMSF=False,
                noStokesI=False, phiNoise_radm2=1e6, nBits=32, showPlots=False,
                debug=False, verbose=False, log=print,units='Jy/beam', prefixOut="prefixOut", args=None):
    """Run RM synthesis on 1D data.

    Args:
        data (list): Contains frequency and polarization data as either:
            [freq_Hz, I, Q, U, dI, dQ, dU]
                freq_Hz (array_like): Frequency of each channel in Hz.
                I (array_like): Stokes I intensity in each channel.
                Q (array_like): Stokes Q intensity in each channel.
                U (array_like): Stokes U intensity in each channel.
                dI (array_like): Error in Stokes I intensity in each channel.
                dQ (array_like): Error in Stokes Q intensity in each channel.
                dU (array_like): Error in Stokes U intensity in each channel.
            or
            [freq_Hz, q, u,  dq, du]
                freq_Hz (array_like): Frequency of each channel in Hz.
                q (array_like): Fractional Stokes Q intensity (Q/I) in each channel.
                u (array_like): Fractional Stokes U intensity (U/I) in each channel.
                dq (array_like): Error in fractional Stokes Q intensity in each channel.
                du (array_like): Error in fractional Stokes U intensity in each channel.

    Kwargs:
        polyOrd (int): Order of polynomial to fit to Stokes I spectrum.
        phiMax_radm2 (float): Maximum absolute Faraday depth (rad/m^2).
        dPhi_radm2 (float): Faraday depth channel size (rad/m^2).
        nSamples (float): Number of samples across the RMSF.
        weightType (str): Can be "variance" or "uniform"
            "variance" -- Weight by uncertainty in Q and U.
            "uniform" -- Weight uniformly (i.e. with 1s)
        fitRMSF (bool): Fit a Gaussian to the RMSF?
        noStokesI (bool: Is Stokes I data provided?
        phiNoise_radm2 (float): ????
        nBits (int): Precision of floating point numbers.
        showPlots (bool): Show plots?
        debug (bool): Turn on debugging messages & plots?
        verbose (bool): Verbosity.
        log (function): Which logging function to use.
        units (str): Units of data.

    Returns:
        mDict (dict): Summary of RM synthesis results.
        aDict (dict): Data output by RM synthesis.

    """
    # Sanity checks
    if not os.path.exists(args.dataFile[0]):
        print("File does not exist: '%s'." % args.dataFile[0])
        sys.exit()
    prefixOut, ext = os.path.splitext(args.dataFile[0])

    # Default data types
    dtFloat = "float" + str(nBits)
    dtComplex = "complex" + str(2*nBits)

    # freq_Hz, I, Q, U, dI, dQ, dU
    try:
        if verbose: log("> Trying [freq_Hz, I, Q, U, dI, dQ, dU]", end=' ')
        (freqArr_Hz, IArr, QArr, UArr, dIArr, dQArr, dUArr) = data
        if verbose: log("... success.")
    except Exception:
        if verbose: log("...failed.")
        # freq_Hz, q, u, dq, du
        try:
            if verbose: log("> Trying [freq_Hz, q, u,  dq, du]", end=' ')
            (freqArr_Hz, QArr, UArr, dQArr, dUArr) = data
            if verbose: log("... success.")
            noStokesI = True
        except Exception:
            if verbose: log("...failed.")
            if debug:
                log(traceback.format_exc())
            sys.exit()
    if verbose: log("Successfully read in the Stokes spectra.")

    # If no Stokes I present, create a dummy spectrum = unity
    if noStokesI:
        if verbose: log("Warn: no Stokes I data in use.")
        IArr = np.ones_like(QArr)
        dIArr = np.zeros_like(QArr)

    # Convert to GHz for convenience
    freqArr_GHz = freqArr_Hz / 1e9
    dQUArr = (dQArr + dUArr)/2.0

    # Fit the Stokes I spectrum and create the fractional spectra
    IModArr, qArr, uArr, dqArr, duArr, fitDict = \
             create_frac_spectra(freqArr  = freqArr_GHz,
                                 IArr     = IArr,
                                 QArr     = QArr,
                                 UArr     = UArr,
                                 dIArr    = dIArr,
                                 dQArr    = dQArr,
                                 dUArr    = dUArr,
                                 polyOrd  = polyOrd,
                                 verbose  = True,
                                 debug    = debug)

    # Plot the data and the Stokes I model fit
    if verbose: log("Plotting the input data and spectral index fit.")
    freqHirArr_Hz =  np.linspace(freqArr_Hz[0], freqArr_Hz[-1], 10000)
    IModHirArr = poly5(fitDict["p"])(freqHirArr_Hz/1e9)
    specFig = plt.figure(figsize=(12.0, 8))
    plot_Ipqu_spectra_fig(freqArr_Hz     = freqArr_Hz,
                          IArr           = IArr,
                          qArr           = qArr,
                          uArr           = uArr,
                          dIArr          = dIArr,
                          dqArr          = dqArr,
                          duArr          = duArr,
                          freqHirArr_Hz  = freqHirArr_Hz,
                          IModArr        = IModHirArr,
                          fig            = specFig,
                          units          = units)

    # Use the custom navigation toolbar (does not work on Mac OS X)
#        try:
#            specFig.canvas.toolbar.pack_forget()
#            CustomNavbar(specFig.canvas, specFig.canvas.toolbar.window)
#        except Exception:
#            pass

    # Display the figure
#        if not plt.isinteractive():
#            specFig.show()

    # DEBUG (plot the Q, U and average RMS spectrum)
    if debug:
        rmsFig = plt.figure(figsize=(12.0, 8))
        ax = rmsFig.add_subplot(111)
        ax.plot(freqArr_Hz/1e9, dQUArr, marker='o', color='k', lw=0.5,
                label='rms <QU>')
        ax.plot(freqArr_Hz/1e9, dQArr, marker='o', color='b', lw=0.5,
                label='rms Q')
        ax.plot(freqArr_Hz/1e9, dUArr, marker='o', color='r', lw=0.5,
                label='rms U')
        xRange = (np.nanmax(freqArr_Hz)-np.nanmin(freqArr_Hz))/1e9
        ax.set_xlim( np.min(freqArr_Hz)/1e9 - xRange*0.05,
                     np.max(freqArr_Hz)/1e9 + xRange*0.05)
        ax.set_xlabel('$\\nu$ (GHz)')
        ax.set_ylabel('RMS '+units)
        ax.set_title("RMS noise in Stokes Q, U and <Q,U> spectra")
#            rmsFig.show()

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

    # Calculate some wavelength parameters
    lambdaSqArr_m2 = np.power(C/freqArr_Hz, 2.0)
    dFreq_Hz = np.nanmin(np.abs(np.diff(freqArr_Hz)))
    lambdaSqRange_m2 = ( np.nanmax(lambdaSqArr_m2) -
                         np.nanmin(lambdaSqArr_m2) )
    dLambdaSqMin_m2 = np.nanmin(np.abs(np.diff(lambdaSqArr_m2)))
    dLambdaSqMax_m2 = np.nanmax(np.abs(np.diff(lambdaSqArr_m2)))

    # Set the Faraday depth range
    fwhmRMSF_radm2 = 2.0 * m.sqrt(3.0) / lambdaSqRange_m2
    if dPhi_radm2 is None:
        dPhi_radm2 = fwhmRMSF_radm2 / nSamples
    if phiMax_radm2 is None:
        phiMax_radm2 = m.sqrt(3.0) / dLambdaSqMax_m2
        phiMax_radm2 = max(phiMax_radm2, fwhmRMSF_radm2*10.)    # Force the minimum phiMax to 10 FWHM

    # Faraday depth sampling. Zero always centred on middle channel
    nChanRM = int(round(abs((phiMax_radm2 - 0.0) / dPhi_radm2)) * 2.0 + 1.0)
    startPhi_radm2 = - (nChanRM-1.0) * dPhi_radm2 / 2.0
    stopPhi_radm2 = + (nChanRM-1.0) * dPhi_radm2 / 2.0
    phiArr_radm2 = np.linspace(startPhi_radm2, stopPhi_radm2, nChanRM)
    phiArr_radm2 = phiArr_radm2.astype(dtFloat)
    if verbose: log("PhiArr = %.2f to %.2f by %.2f (%d chans)." % (phiArr_radm2[0],
                                                         phiArr_radm2[-1],
                                                         float(dPhi_radm2),
                                                         nChanRM))

    # Calculate the weighting as 1/sigma^2 or all 1s (uniform)
    if weightType=="variance":
        weightArr = 1.0 / np.power(dQUArr, 2.0)
    else:
        weightType = "uniform"
        weightArr = np.ones(freqArr_Hz.shape, dtype=dtFloat)
    if verbose: log("Weight type is '%s'." % weightType)

    startTime = time.time()

    # Perform RM-synthesis on the spectrum
    dirtyFDF, lam0Sq_m2 = do_rmsynth_planes(dataQ           = qArr,
                                            dataU           = uArr,
                                            lambdaSqArr_m2  = lambdaSqArr_m2,
                                            phiArr_radm2    = phiArr_radm2,
                                            weightArr       = weightArr,
                                            nBits           = nBits,
                                            verbose         = verbose,
                                            log             = log)

    # Calculate the Rotation Measure Spread Function
    RMSFArr, phi2Arr_radm2, fwhmRMSFArr, fitStatArr = \
        get_rmsf_planes(lambdaSqArr_m2  = lambdaSqArr_m2,
                        phiArr_radm2    = phiArr_radm2,
                        weightArr       = weightArr,
                        mskArr          = ~np.isfinite(qArr),
                        lam0Sq_m2       = lam0Sq_m2,
                        double          = True,
                        fitRMSF         = fitRMSF,
                        fitRMSFreal     = False,
                        nBits           = nBits,
                        verbose         = verbose,
                        log             = log)
    fwhmRMSF = float(fwhmRMSFArr)

    # ALTERNATE RM-SYNTHESIS CODE --------------------------------------------#

    #dirtyFDF, [phi2Arr_radm2, RMSFArr], lam0Sq_m2, fwhmRMSF = \
    #          do_rmsynth(qArr, uArr, lambdaSqArr_m2, phiArr_radm2, weightArr)

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

    endTime = time.time()
    cputime = (endTime - startTime)
    if verbose: log("> RM-synthesis completed in %.2f seconds." % cputime)

    # Determine the Stokes I value at lam0Sq_m2 from the Stokes I model
    # Multiply the dirty FDF by Ifreq0 to recover the PI
    freq0_Hz = C / m.sqrt(lam0Sq_m2)
    Ifreq0 = poly5(fitDict["p"])(freq0_Hz/1e9)
    dirtyFDF *= (Ifreq0)    # FDF is in fracpol units initially, convert back to flux

    # Calculate the theoretical noise in the FDF !!Old formula only works for wariance weights!
    weightArr = np.where(np.isnan(weightArr), 0.0, weightArr)
    dFDFth = np.sqrt( np.sum(weightArr**2 * np.nan_to_num(dQUArr)**2) / (np.sum(weightArr))**2 )


    # Measure the parameters of the dirty FDF
    # Use the theoretical noise to calculate uncertainties
    mDict = measure_FDF_parms(FDF         = dirtyFDF,
                              phiArr      = phiArr_radm2,
                              fwhmRMSF    = fwhmRMSF,
                              dFDF        = dFDFth,
                              lamSqArr_m2 = lambdaSqArr_m2,
                              lam0Sq      = lam0Sq_m2)
    mDict["Ifreq0"] = toscalar(Ifreq0)
    mDict["polyCoeffs"] =  ",".join([str(x) for x in fitDict["p"]])
    mDict["IfitStat"] = fitDict["fitStatus"]
    mDict["IfitChiSqRed"] = fitDict["chiSqRed"]
    mDict["lam0Sq_m2"] = toscalar(lam0Sq_m2)
    mDict["freq0_Hz"] = toscalar(freq0_Hz)
    mDict["fwhmRMSF"] = toscalar(fwhmRMSF)
    mDict["dQU"] = toscalar(nanmedian(dQUArr))
    mDict["dFDFth"] = toscalar(dFDFth)
    mDict["units"] = units
    
    if fitDict["fitStatus"] >= 128:
        log("WARNING: Stokes I model contains negative values!")
    elif fitDict["fitStatus"] >= 64:
        log("Caution: Stokes I model has low signal-to-noise.")



    #Add information on nature of channels:
    good_channels=np.where(np.logical_and(weightArr != 0,np.isfinite(qArr)))[0]
    mDict["min_freq"]=float(np.min(freqArr_Hz[good_channels]))
    mDict["max_freq"]=float(np.max(freqArr_Hz[good_channels]))
    mDict["N_channels"]=good_channels.size
    mDict["median_channel_width"]=float(np.median(np.diff(freqArr_Hz)))

    # Measure the complexity of the q and u spectra
    mDict["fracPol"] = mDict["ampPeakPIfit"]/(Ifreq0)
    mD, pD = measure_qu_complexity(freqArr_Hz = freqArr_Hz,
                                   qArr       = qArr,
                                   uArr       = uArr,
                                   dqArr      = dqArr,
                                   duArr      = duArr,
                                   fracPol    = mDict["fracPol"],
                                   psi0_deg   = mDict["polAngle0Fit_deg"],
                                   RM_radm2   = mDict["phiPeakPIfit_rm2"])
    mDict.update(mD)

    # Debugging plots for spectral complexity measure
    if debug:
        tmpFig = plot_complexity_fig(xArr=pD["xArrQ"],
                                     qArr=pD["yArrQ"],
                                     dqArr=pD["dyArrQ"],
                                     sigmaAddqArr=pD["sigmaAddArrQ"],
                                     chiSqRedqArr=pD["chiSqRedArrQ"],
                                     probqArr=pD["probArrQ"],
                                     uArr=pD["yArrU"],
                                     duArr=pD["dyArrU"],
                                     sigmaAdduArr=pD["sigmaAddArrU"],
                                     chiSqReduArr=pD["chiSqRedArrU"],
                                     probuArr=pD["probArrU"],
                                     mDict=mDict)
        if saveOutput:
            if verbose: print("Saving debug plots:")
            outFilePlot = prefixOut + ".debug-plots.pdf"
            if verbose: print("> " + outFilePlot)
            tmpFig.savefig(outFilePlot, bbox_inches = 'tight')
        else:
            tmpFig.show()

    #add array dictionary
    aDict = dict()
    aDict["phiArr_radm2"] = phiArr_radm2
    aDict["phi2Arr_radm2"] = phi2Arr_radm2
    aDict["RMSFArr"] = RMSFArr
    aDict["freqArr_Hz"] = freqArr_Hz
    aDict["weightArr"]=weightArr
    aDict["dirtyFDF"]=dirtyFDF

    if verbose:
       # Print the results to the screen
       log()
       log('-'*80)
       log('RESULTS:\n')
       log('FWHM RMSF = %.4g rad/m^2' % (mDict["fwhmRMSF"]))

       log('Pol Angle = %.4g (+/-%.4g) deg' % (mDict["polAngleFit_deg"],
                                              mDict["dPolAngleFit_deg"]))
       log('Pol Angle 0 = %.4g (+/-%.4g) deg' % (mDict["polAngle0Fit_deg"],
                                                mDict["dPolAngle0Fit_deg"]))
       log('Peak FD = %.4g (+/-%.4g) rad/m^2' % (mDict["phiPeakPIfit_rm2"],
                                                mDict["dPhiPeakPIfit_rm2"]))
       log('freq0_GHz = %.4g ' % (mDict["freq0_Hz"]/1e9))
       log('I freq0 = %.4g %s' % (mDict["Ifreq0"],units))
       log('Peak PI = %.4g (+/-%.4g) %s' % (mDict["ampPeakPIfit"],
                                                mDict["dAmpPeakPIfit"],units))
       log('QU Noise = %.4g %s' % (mDict["dQU"],units))
       log('FDF Noise (theory)   = %.4g %s' % (mDict["dFDFth"],units))
       log('FDF Noise (Corrected MAD) = %.4g %s' % (mDict["dFDFcorMAD"],units))
       log('FDF Noise (rms)   = %.4g %s' % (mDict["dFDFrms"],units))
       log('FDF SNR = %.4g ' % (mDict["snrPIfit"]))
       log('sigma_add(q) = %.4g (+%.4g, -%.4g)' % (mDict["sigmaAddQ"],
                                            mDict["dSigmaAddPlusQ"],
                                            mDict["dSigmaAddMinusQ"]))
       log('sigma_add(u) = %.4g (+%.4g, -%.4g)' % (mDict["sigmaAddU"],
                                            mDict["dSigmaAddPlusU"],
                                            mDict["dSigmaAddMinusU"]))
       log()
       log('-'*80)



    # Plot the RM Spread Function and dirty FDF
    if showPlots or saveOutput:
        fdfFig = plt.figure(figsize=(12.0, 8))
        plot_rmsf_fdf_fig(phiArr     = phiArr_radm2,
                          FDF        = dirtyFDF,
                          phi2Arr    = phi2Arr_radm2,
                          RMSFArr    = RMSFArr,
                          fwhmRMSF   = fwhmRMSF,
                          vLine      = mDict["phiPeakPIfit_rm2"],
                          fig        = fdfFig,
                          units      = units)

        # Use the custom navigation toolbar
#        try:
#            fdfFig.canvas.toolbar.pack_forget()
#            CustomNavbar(fdfFig.canvas, fdfFig.canvas.toolbar.window)
#        except Exception:
#            pass

        # Display the figure
#        fdfFig.show()

    # Pause if plotting enabled
    if showPlots:
        plt.show()
    elif saveOutput or debug:
        if verbose: print("Saving RMSF and dirty FDF plot:")
        outFilePlot = prefixOut + ".RMSF-dirtyFDF-plots.pdf"
        if verbose: print("> " + outFilePlot)
        fdfFig.savefig(outFilePlot, bbox_inches = 'tight')
        #        #if verbose: print "Press <RETURN> to exit ...",
#        input()

    return mDict, aDict

def readFile(dataFile, nBits, verbose=True, debug=False):
    """
    Read the I, Q & U data from the ASCII file.
    
    Inputs:
        datafile (str): relative or absolute path to file.
        nBits (int): number of bits to store the data as.
        verbose (bool): Print verbose messages to terminal?
        debug (bool): Print full traceback in case of failure?
        
    Returns:
        data (list of arrays): List containing the columns found in the file.
        If Stokes I is present, this will be [freq_Hz, I, Q, U, dI, dQ, dU], 
        else [freq_Hz, q, u,  dq, du].
    """

    # Default data types
    dtFloat = "float" + str(nBits)
    dtComplex = "complex" + str(2*nBits)

    # Output prefix is derived from the input file name


    # Read the data-file. Format=space-delimited, comments="#".
    if verbose: print("Reading the data file '%s':" % dataFile)
    # freq_Hz, I, Q, U, dI, dQ, dU
    try:
        if verbose: print("> Trying [freq_Hz, I, Q, U, dI, dQ, dU]", end=' ')
        (freqArr_Hz, IArr, QArr, UArr,
         dIArr, dQArr, dUArr) = \
         np.loadtxt(dataFile, unpack=True, dtype=dtFloat)
        if verbose: print("... success.")
        data=[freqArr_Hz, IArr, QArr, UArr, dIArr, dQArr, dUArr]
    except Exception:
        if verbose: print("...failed.")
        # freq_Hz, q, u, dq, du
        try:
            if verbose: print("> Trying [freq_Hz, q, u,  dq, du]", end=' ')
            (freqArr_Hz, QArr, UArr, dQArr, dUArr) = \
                         np.loadtxt(dataFile, unpack=True, dtype=dtFloat)
            if verbose: print("... success.")
            data=[freqArr_Hz, QArr, UArr, dQArr, dUArr]

            noStokesI = True
        except Exception:
            if verbose: print("...failed.")
            if debug:
                print(traceback.format_exc())
            sys.exit()
    if verbose: print("Successfully read in the Stokes spectra.")
    return data

def saveOutput(outdict, arrdict, prefixOut, verbose):
    # Save the  dirty FDF, RMSF and weight array to ASCII files
    if verbose: print("Saving the dirty FDF, RMSF weight arrays to ASCII files.")
    outFile = prefixOut + "_FDFdirty.dat"
    if verbose:
        print("> %s" % outFile)
    np.savetxt(outFile, list(zip(arrdict["phiArr_radm2"], arrdict["dirtyFDF"].real, arrdict["dirtyFDF"].imag)))

    outFile = prefixOut + "_RMSF.dat"
    if verbose:
        print("> %s" % outFile)
    np.savetxt(outFile, list(zip(arrdict["phi2Arr_radm2"], arrdict["RMSFArr"].real, arrdict["RMSFArr"].imag)))

    outFile = prefixOut + "_weight.dat"
    if verbose:
        print("> %s" % outFile)
    np.savetxt(outFile, list(zip(arrdict["freqArr_Hz"], arrdict["weightArr"])))

    # Save the measurements to a "key=value" text file
    outFile = prefixOut + "_RMsynth.dat"

    if verbose:
        print("Saving the measurements on the FDF in 'key=val' and JSON formats.")
        print("> %s" % outFile)

    FH = open(outFile, "w")
    for k, v in outdict.items():
        FH.write("%s=%s\n" % (k, v))
    FH.close()


    outFile = prefixOut + "_RMsynth.json"

    if verbose:
        print("> %s" % outFile)
    json.dump(dict(outdict), open(outFile, "w"))


#-----------------------------------------------------------------------------#
def main():
    import argparse
    """
    Start the function to perform RM-synthesis if called from the command line.
    """

    # Help string to be shown using the -h option
    descStr = """
    Run RM-synthesis on Stokes I, Q and U spectra (1D) stored in an ASCII
    file. The Stokes I spectrum is first fit with a polynomial and the
    resulting model used to create fractional q = Q/I and u = U/I spectra.

    The ASCII file should the following columns, in a space separated format:
    [freq_Hz, I, Q, U, I_err, Q_err, U_err]
    OR
    [freq_Hz, Q, U, Q_err, U_err]


    To get outputs, one or more of the following flags must be set: -S, -p, -v.
    """

    epilog_text="""
    Outputs with -S flag:
    _FDFdirty.dat: Dirty FDF/RM Spectrum [Phi, Q, U]
    _RMSF.dat: Computed RMSF [Phi, Q, U]
    _RMsynth.dat: list of derived parameters for RM spectrum
                (approximately equivalent to -v flag output)
    _RMsynth.json: dictionary of derived parameters for RM spectrum
    _weight.dat: Calculated channel weights [freq_Hz, weight]
    """

    # Parse the command line options
    parser = argparse.ArgumentParser(description=descStr,epilog=epilog_text,
                                 formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("dataFile", metavar="dataFile.dat", nargs=1,
                        help="ASCII file containing Stokes spectra & errors.")
    parser.add_argument("-t", dest="fitRMSF", action="store_true",
                        help="fit a Gaussian to the RMSF [False]")
    parser.add_argument("-l", dest="phiMax_radm2", type=float, default=None,
                        help="absolute max Faraday depth sampled [Auto].")
    parser.add_argument("-d", dest="dPhi_radm2", type=float, default=None,
                        help="width of Faraday depth channel [Auto].\n(overrides -s NSAMPLES flag)")
    parser.add_argument("-s", dest="nSamples", type=float, default=10,
                        help="number of samples across the RMSF lobe [10].")
    parser.add_argument("-w", dest="weightType", default="variance",
                        help="weighting [inverse variance] or 'uniform' (all 1s).")
    parser.add_argument("-o", dest="polyOrd", type=int, default=2,
                        help="polynomial order to fit to I spectrum [2].")
    parser.add_argument("-i", dest="noStokesI", action="store_true",
                        help="ignore the Stokes I spectrum [False].")
    parser.add_argument("-b", dest="bit64", action="store_true",
                        help="use 64-bit floating point precision [False (uses 32-bit)]")
    parser.add_argument("-p", dest="showPlots", action="store_true",
                        help="show the plots [False].")
    parser.add_argument("-v", dest="verbose", action="store_true",
                        help="verbose output [False].")
    parser.add_argument("-S", dest="saveOutput", action="store_true",
                        help="save the arrays and plots [False].")
    parser.add_argument("-D", dest="debug", action="store_true",
                        help="turn on debugging messages & plots [False].")
    parser.add_argument("-U", dest="units", type=str, default="Jy/beam",
                        help="Intensity units of the data. [Jy/beam]")
    args = parser.parse_args()

    # Sanity checks
    if not os.path.exists(args.dataFile[0]):
        print("File does not exist: '%s'." % args.dataFile[0])
        sys.exit()
    prefixOut, ext = os.path.splitext(args.dataFile[0])
    dataDir, dummy = os.path.split(args.dataFile[0])
    # Set the floating point precision
    nBits = 32
    if args.bit64:
        nBits = 64
    verbose=args.verbose
    data = readFile(args.dataFile[0],nBits, verbose=verbose, debug=args.debug)

    # Run RM-synthesis on the spectra
    mDict, aDict = run_rmsynth(data           = data,
                polyOrd        = args.polyOrd,
                phiMax_radm2   = args.phiMax_radm2,
                dPhi_radm2     = args.dPhi_radm2,
                nSamples       = args.nSamples,
                weightType     = args.weightType,
                fitRMSF        = args.fitRMSF,
                noStokesI      = args.noStokesI,
                nBits          = nBits,
                showPlots      = args.showPlots,
                debug          = args.debug,
                verbose        = verbose,
                units          = args.units,
                prefixOut      = prefixOut,
                args           = args,
                )

    if args.saveOutput:
        saveOutput(mDict, aDict, prefixOut, verbose)


#-----------------------------------------------------------------------------#
if __name__ == "__main__":
    main()

[((121, 21, 121, 55), 'os.path.splitext', 'os.path.splitext', ({(121, 38, 121, 54): 'args.dataFile[0]'}, {}), '(args.dataFile[0])', False, 'import os\n'), ((159, 13, 168, 50), 'RMutils.util_misc.create_frac_spectra', 'create_frac_spectra', (), '', False, 'from RMutils.util_misc import create_frac_spectra\n'), ((172, 21, 172, 70), 'numpy.linspace', 'np.linspace', ({(172, 33, 172, 46): 'freqArr_Hz[0]', (172, 48, 172, 62): 'freqArr_Hz[-1]', (172, 64, 172, 69): '10000'}, {}), '(freqArr_Hz[0], freqArr_Hz[-1], 10000)', True, 'import numpy as np\n'), ((174, 14, 174, 43), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((175, 4, 185, 49), 'RMutils.util_plotTk.plot_Ipqu_spectra_fig', 'plot_Ipqu_spectra_fig', (), '', False, 'from RMutils.util_plotTk import plot_Ipqu_spectra_fig\n'), ((219, 21, 219, 48), 'numpy.power', 'np.power', ({(219, 30, 219, 42): 'C / freqArr_Hz', (219, 44, 219, 47): '2.0'}, {}), '(C / freqArr_Hz, 2.0)', True, 'import numpy as np\n'), ((238, 19, 238, 70), 'numpy.linspace', 'np.linspace', ({(238, 31, 238, 45): 'startPhi_radm2', (238, 47, 238, 60): 'stopPhi_radm2', (238, 62, 238, 69): 'nChanRM'}, {}), '(startPhi_radm2, stopPhi_radm2, nChanRM)', True, 'import numpy as np\n'), ((253, 16, 253, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((256, 26, 263, 66), 'RMutils.util_RM.do_rmsynth_planes', 'do_rmsynth_planes', (), '', False, 'from RMutils.util_RM import do_rmsynth_planes\n'), ((287, 14, 287, 25), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((304, 12, 309, 54), 'RMutils.util_RM.measure_FDF_parms', 'measure_FDF_parms', (), '', False, 'from RMutils.util_RM import measure_FDF_parms\n'), ((310, 22, 310, 38), 'RMutils.util_misc.toscalar', 'toscalar', ({(310, 31, 310, 37): 'Ifreq0'}, {}), '(Ifreq0)', False, 'from RMutils.util_misc import toscalar\n'), ((314, 25, 314, 44), 'RMutils.util_misc.toscalar', 'toscalar', ({(314, 34, 314, 43): 'lam0Sq_m2'}, {}), '(lam0Sq_m2)', False, 'from RMutils.util_misc import toscalar\n'), ((315, 24, 315, 42), 'RMutils.util_misc.toscalar', 'toscalar', ({(315, 33, 315, 41): 'freq0_Hz'}, {}), '(freq0_Hz)', False, 'from RMutils.util_misc import toscalar\n'), ((316, 24, 316, 42), 'RMutils.util_misc.toscalar', 'toscalar', ({(316, 33, 316, 41): 'fwhmRMSF'}, {}), '(fwhmRMSF)', False, 'from RMutils.util_misc import toscalar\n'), ((318, 22, 318, 38), 'RMutils.util_misc.toscalar', 'toscalar', ({(318, 31, 318, 37): 'dFDFth'}, {}), '(dFDFth)', False, 'from RMutils.util_misc import toscalar\n'), ((337, 13, 344, 74), 'RMutils.util_RM.measure_qu_complexity', 'measure_qu_complexity', (), '', False, 'from RMutils.util_RM import measure_qu_complexity\n'), ((569, 13, 570, 79), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((605, 21, 605, 55), 'os.path.splitext', 'os.path.splitext', ({(605, 38, 605, 54): 'args.dataFile[0]'}, {}), '(args.dataFile[0])', False, 'import os\n'), ((606, 21, 606, 52), 'os.path.split', 'os.path.split', ({(606, 35, 606, 51): 'args.dataFile[0]'}, {}), '(args.dataFile[0])', False, 'import os\n'), ((118, 11, 118, 43), 'os.path.exists', 'os.path.exists', ({(118, 26, 118, 42): 'args.dataFile[0]'}, {}), '(args.dataFile[0])', False, 'import os\n'), ((120, 8, 120, 18), 'sys.exit', 'sys.exit', ({}, {}), '()', False, 'import sys\n'), ((150, 15, 150, 33), 'numpy.ones_like', 'np.ones_like', ({(150, 28, 150, 32): 'QArr'}, {}), '(QArr)', True, 'import numpy as np\n'), ((151, 16, 151, 35), 'numpy.zeros_like', 'np.zeros_like', ({(151, 30, 151, 34): 'QArr'}, {}), '(QArr)', True, 'import numpy as np\n'), ((173, 17, 173, 36), 'RMutils.util_misc.poly5', 'poly5', ({(173, 23, 173, 35): "fitDict['p']"}, {}), "(fitDict['p'])", False, 'from RMutils.util_misc import poly5\n'), ((200, 17, 200, 46), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((221, 25, 221, 50), 'numpy.nanmax', 'np.nanmax', ({(221, 35, 221, 49): 'lambdaSqArr_m2'}, {}), '(lambdaSqArr_m2)', True, 'import numpy as np\n'), ((222, 25, 222, 50), 'numpy.nanmin', 'np.nanmin', ({(222, 35, 222, 49): 'lambdaSqArr_m2'}, {}), '(lambdaSqArr_m2)', True, 'import numpy as np\n'), ((250, 20, 250, 60), 'numpy.ones', 'np.ones', (), '', True, 'import numpy as np\n'), ((293, 19, 293, 36), 'math.sqrt', 'm.sqrt', ({(293, 26, 293, 35): 'lam0Sq_m2'}, {}), '(lam0Sq_m2)', True, 'import math as m\n'), ((294, 13, 294, 32), 'RMutils.util_misc.poly5', 'poly5', ({(294, 19, 294, 31): "fitDict['p']"}, {}), "(fitDict['p'])", False, 'from RMutils.util_misc import poly5\n'), ((298, 25, 298, 44), 'numpy.isnan', 'np.isnan', ({(298, 34, 298, 43): 'weightArr'}, {}), '(weightArr)', True, 'import numpy as np\n'), ((317, 28, 317, 45), 'RMutils.util_misc.nanmedian', 'nanmedian', ({(317, 38, 317, 44): 'dQUArr'}, {}), '(dQUArr)', False, 'from RMutils.util_misc import nanmedian\n'), ((330, 28, 330, 61), 'numpy.min', 'np.min', ({(330, 35, 330, 60): 'freqArr_Hz[good_channels]'}, {}), '(freqArr_Hz[good_channels])', True, 'import numpy as np\n'), ((331, 28, 331, 61), 'numpy.max', 'np.max', ({(331, 35, 331, 60): 'freqArr_Hz[good_channels]'}, {}), '(freqArr_Hz[good_channels])', True, 'import numpy as np\n'), ((349, 17, 360, 49), 'RMutils.util_plotTk.plot_complexity_fig', 'plot_complexity_fig', (), '', False, 'from RMutils.util_plotTk import plot_complexity_fig\n'), ((413, 17, 413, 46), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((414, 8, 421, 45), 'RMutils.util_plotTk.plot_rmsf_fdf_fig', 'plot_rmsf_fdf_fig', (), '', False, 'from RMutils.util_plotTk import plot_rmsf_fdf_fig\n'), ((435, 8, 435, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((476, 9, 476, 57), 'numpy.loadtxt', 'np.loadtxt', (), '', True, 'import numpy as np\n'), ((602, 11, 602, 43), 'os.path.exists', 'os.path.exists', ({(602, 26, 602, 42): 'args.dataFile[0]'}, {}), '(args.dataFile[0])', False, 'import os\n'), ((604, 8, 604, 18), 'sys.exit', 'sys.exit', ({}, {}), '()', False, 'import sys\n'), ((220, 32, 220, 51), 'numpy.diff', 'np.diff', ({(220, 40, 220, 50): 'freqArr_Hz'}, {}), '(freqArr_Hz)', True, 'import numpy as np\n'), ((223, 39, 223, 62), 'numpy.diff', 'np.diff', ({(223, 47, 223, 61): 'lambdaSqArr_m2'}, {}), '(lambdaSqArr_m2)', True, 'import numpy as np\n'), ((224, 39, 224, 62), 'numpy.diff', 'np.diff', ({(224, 47, 224, 61): 'lambdaSqArr_m2'}, {}), '(lambdaSqArr_m2)', True, 'import numpy as np\n'), ((227, 27, 227, 38), 'math.sqrt', 'm.sqrt', ({(227, 34, 227, 37): '(3.0)'}, {}), '(3.0)', True, 'import math as m\n'), ((231, 23, 231, 34), 'math.sqrt', 'm.sqrt', ({(231, 30, 231, 33): '(3.0)'}, {}), '(3.0)', True, 'import math as m\n'), ((247, 26, 247, 47), 'numpy.power', 'np.power', ({(247, 35, 247, 41): 'dQUArr', (247, 43, 247, 46): '(2.0)'}, {}), '(dQUArr, 2.0)', True, 'import numpy as np\n'), ((333, 50, 333, 69), 'numpy.diff', 'np.diff', ({(333, 58, 333, 68): 'freqArr_Hz'}, {}), '(freqArr_Hz)', True, 'import numpy as np\n'), ((208, 18, 208, 39), 'numpy.nanmax', 'np.nanmax', ({(208, 28, 208, 38): 'freqArr_Hz'}, {}), '(freqArr_Hz)', True, 'import numpy as np\n'), ((208, 40, 208, 61), 'numpy.nanmin', 'np.nanmin', ({(208, 50, 208, 60): 'freqArr_Hz'}, {}), '(freqArr_Hz)', True, 'import numpy as np\n'), ((270, 43, 270, 60), 'numpy.isfinite', 'np.isfinite', ({(270, 55, 270, 59): 'qArr'}, {}), '(qArr)', True, 'import numpy as np\n'), ((299, 73, 299, 90), 'numpy.sum', 'np.sum', ({(299, 80, 299, 89): 'weightArr'}, {}), '(weightArr)', True, 'import numpy as np\n'), ((329, 57, 329, 74), 'numpy.isfinite', 'np.isfinite', ({(329, 69, 329, 73): 'qArr'}, {}), '(qArr)', True, 'import numpy as np\n'), ((485, 25, 485, 73), 'numpy.loadtxt', 'np.loadtxt', (), '', True, 'import numpy as np\n'), ((144, 12, 144, 22), 'sys.exit', 'sys.exit', ({}, {}), '()', False, 'import sys\n'), ((209, 21, 209, 39), 'numpy.min', 'np.min', ({(209, 28, 209, 38): 'freqArr_Hz'}, {}), '(freqArr_Hz)', True, 'import numpy as np\n'), ((210, 21, 210, 39), 'numpy.max', 'np.max', ({(210, 28, 210, 38): 'freqArr_Hz'}, {}), '(freqArr_Hz)', True, 'import numpy as np\n'), ((494, 12, 494, 22), 'sys.exit', 'sys.exit', ({}, {}), '()', False, 'import sys\n'), ((299, 44, 299, 65), 'numpy.nan_to_num', 'np.nan_to_num', ({(299, 58, 299, 64): 'dQUArr'}, {}), '(dQUArr)', True, 'import numpy as np\n'), ((143, 20, 143, 42), 'traceback.format_exc', 'traceback.format_exc', ({}, {}), '()', False, 'import traceback\n'), ((493, 22, 493, 44), 'traceback.format_exc', 'traceback.format_exc', ({}, {}), '()', False, 'import traceback\n')]