#Imports
import sys, os, math
from AIPS import AIPS, AIPSDisk
from AIPSTask import AIPSTask, AIPSList
from AIPSData import AIPSUVData, AIPSImage, AIPSCat
from optparse import OptionParser

#Options parser
usage = "usage: %prog -e <experiment name> -t <targets> [options]"
parser = OptionParser(usage)
parser.add_option("-e", "--experiment", dest="experiment", default="",
                  help="experiment name eg v190x", metavar="EXPERIMENT")
parser.add_option("-t", "--targets", dest="targets", default="",
                  help="Comma separated list of pulsars",
                  metavar="target1[,target2,...targetN]")
parser.add_option("-p", "--prefix", dest="prefix", default="junk",
                  help="prefix for uv filenames")
parser.add_option("-s", "--suffix", dest="suffix", default=".uvf",
                  help="suffix for uv filenames")
parser.add_option("-o", "--outputsuffix", dest="outsuffix", default=".recal",
                  help="suffix for jmfitstats filenames")
parser.add_option("-r", "--reweight", dest="reweight", default=False,
                  action="store_true", help="Uvaver with `180,true")
parser.add_option("-w", "--weightstring", dest="weightstring", default="2,-1",
                  help="Weighting string for uvweight in difmap")
parser.add_option("-k", "--killweights", dest="killweights",
                  action="store_true", default=False,
                  help="Set all weights to 1 when loading (default False)")
parser.add_option("--threshold", dest="threshold", default="-1.0",
                  help="% Threshold: all below threshold set to 0, all " + \
                  "weights above threshold to 1 when loading (< 0 == off)")
parser.add_option("--trim", dest="trim", action="store_true", default=False, \
                  help="Throw away N/(N+1)% of visibilities")
parser.add_option("-a", "--addnoise", dest="addnoise", action="store_true", \
                  default=False, help="Add thermal noise to all baselines")

(options, junk) = parser.parse_args()

#Global variables
experiment = options.experiment
targets = options.targets.split(',')
uvprefix = options.prefix
uvsuffix = options.suffix
jmsuffix = options.outsuffix
uvweightstr = options.weightstring
killweights = options.killweights
reweight = options.reweight
threshold = float(options.threshold)
dotrim = options.trim
doaddnoise = options.addnoise
dothreshold = not(threshold < 0.0)
rootdir = os.path.dirname('/nfs/cluster/ska/adeller/v190/')
aips_version_std = '31DEC05'
AIPS.userno = 715
difmappath = 'difmap'
if killweights:
    if dotrim:
        difmappath = '/home/ssi/adeller/difmap/dropandunweight_difmap/difmap'
    else:
        difmappath = '/home/ssi/adeller/difmap/noweight_difmap/difmap'
else:
    if dotrim:
        difmappath = '/home/ssi/adeller/difmap/dropsome_difmap/difmap'
if dothreshold:
    difmappath = '/home/ssi/adeller/difmap/weightclip_difmap/difmap'
if doaddnoise:
    difmappath = '/home/ssi/adeller/difmap/addnoise_difmap/difmap'
    difmappath2 = '/home/ssi/adeller/difmap/noweight_difmap/difmap'
isll = True
if experiment == 'v190a' or experiment == 'v190b' or experiment == 'v190e' \
       or experiment == 'v190g' or experiment == 'v190j':
    isll = False
    
#Checks
if len(junk) > 0:
    print "You can only supply the allowed options!!! Aborting"
    sys.exit(1)

if not (uvweightstr == "2,0" or uvweightstr == "2,-1" or  \
        uvweightstr == "0" or uvweightstr == "2,-0.5" or \
        uvweightstr == "0,-1"):
    print "Disallowed uvweight option " + uvweightstr + " - aborting!!!"
    sys.exit()

if not (experiment == "v190a" or experiment == "v190b" or \
        experiment == "v190c" or experiment == "v190d" or \
        experiment == "v190e" or experiment == "v190f" or \
        experiment == "v190g" or experiment == "v190h" or \
        experiment == "v190i" or experiment == "v190j" or \
        experiment == "v190k" or experiment == "v190l"):
    print "Disallowed experiment " + experiment + " - aborting!!!"
    sys.exit()

for pulsar in targets:
    if not (pulsar == "0108-1431" or pulsar == "0437-4715" or \
            pulsar == "0630-2834" or pulsar == "0737-3039" or \
            pulsar == "1559-4438" or pulsar == "2048-1616" or \
            pulsar == "2144-3933" or pulsar == "2145-0750"):
        print "Disallowed pulsar " + pulsar + " - aborting!!!"
        sys.exit()

# GET_RA_DEC
def get_ra_dec(pulsar, uvprefix, uvsuffix):
    uvfitsin = open(rootdir + '/' + experiment + '/' + uvprefix + \
                    pulsar + uvsuffix, 'r')
    keyword = " "*80
    while not keyword[11:13] == 'RA' and len(keyword) == 80:
        keyword = uvfitsin.read(80)
    if not len(keyword) == 80:
        print "Couldn't find RA!!! Aborting."
        sys.exit(1)
    while not keyword[:5] == "CRVAL" and len(keyword) == 80:
        keyword = uvfitsin.read(80)
    if not len(keyword) == 80:
        print "Couldn't find RA!!! Aborting."
        sys.exit(1)
    ra = float(keyword.split()[2])*math.pi/180.0
    while not keyword[11:14] == 'DEC' and len(keyword) == 80:
        keyword = uvfitsin.read(80)
    if not len(keyword) == 80:
        print "Couldn't find Dec!!! Aborting."
        sys.exit(1)
    while not keyword[:5] == "CRVAL" and len(keyword) == 80:
        keyword = uvfitsin.read(80)
    if not len(keyword) == 80:
        print "Couldn't find Dec!!! Aborting."
        sys.exit(1)
    dec = float(keyword.split()[2])*math.pi/180.0
    return (ra, dec)

# DIFMAP_MAPPSR
def difmap_mappsr(source, isll, centre_ra_deg, centre_dec_deg, uvweightstr,
                  experiment, difmappath, uvprefix, uvsuffix, jmsuffix, \
                  saveagain):
    sourceuvfile = rootdir + "/" + experiment + "/" + uvprefix + \
                   source + uvsuffix
    if experiment == 'v190a' or experiment == 'v190e' or \
           experiment == 'v190g' or experiment == 'v190k':
        pixsize = 0.4
    else:
        pixsize = 2.5
    difmapout = open("difmap_commands", "w")
    difmapout.write("float pkflux\n")
    difmapout.write("float peakx\n")
    difmapout.write("float peaky\n")
    difmapout.write("float finepeakx\n")
    difmapout.write("float finepeaky\n")
    difmapout.write("float rmsflux\n")
    difmapout.write("integer ilevs\n")
    difmapout.write("float lowlev\n")
    difmapout.write("obs " + sourceuvfile + "\n")
    difmapout.write("mapsize 1024," + str(pixsize) + "\n")
    if dothreshold:
        if isll:
            difmapout.write("select ll,1,2,3,4\n")
        else:
            difmapout.write("select i,1,2\n")
        if experiment == 'v190k':
            difmapout.write("select rr,1,2,3,4\n")
        difmapout.write("invert\n");
        difmapout.write("wtscale " + str(threshold) + "\n")
        difmapout.write("obs " + sourceuvfile + "\n")
    if reweight:
        difmapout.write("uvaver 180,true\n")
    else:
        difmapout.write("uvaver 20\n")
    difmapout.write("mapcolor none\n")
    difmapout.write("uvweight " + uvweightstr + "\n")

    #Find the peak from the combined first
    if isll:
        difmapout.write("select ll,1,2,3,4\n")
    else:
        difmapout.write("select i,1,2\n")
    if experiment == 'v190k':
        difmapout.write("select rr,1,2,3,4\n")
    difmapout.write("peakx = peak(x,max)\n")
    difmapout.write("peaky = peak(y,max)\n")
    difmapout.write("shift -peakx,-peaky\n")
    difmapout.write("mapsize 1024,0.1\n")
    difmapout.write("finepeakx = peak(x,max)\n")
    difmapout.write("finepeaky = peak(y,max)\n")
    
    #Do each individual band, one at a time
    maxif = 4
    pols = ['rr','ll']
    for i in range(maxif):
        for p in pols:
            difmapout.write("select " + p + "," + str(i+1) + "\n")
            write_difmappsrscript(source, p + "." + str(i+1), difmapout, \
                                  pixsize, jmsuffix)
    #Now the combined one
    if experiment == 'v190k':
        difmapout.write("select rr,1,2,3,4\n")
    elif isll:
        difmapout.write("select ll,1,2,3,4\n")
    else:
        difmapout.write("select i,1,2\n")
    write_difmappsrscript(source, 'combined', difmapout, \
                          pixsize, jmsuffix)
    if saveagain:
        difmapout.write("wobs " + rootdir + "/" + experiment + "/noise" + \
                        source + ".uvf\n")
    difmapout.write("exit\n")

    difmapout.close()
    os.system(difmappath + " < difmap_commands")

# WRITE_DIFMAPPSRSCRIPT
def write_difmappsrscript(source, bands, difmapout, pixsize, jmsuffix):
    difmapout.write("clrmod true\n")
    difmapout.write("unshift\n")
    difmapout.write("shift -peakx,-peaky\n")
    difmapout.write("mapsize 1024,0.1\n")
    difmapout.write("pkflux = peak(flux,max)\n")
    difmapout.write("addcmp pkflux, true, finepeakx, finepeaky, true, 0, " + \
                    "false, 1, false, 0, false, 0, 0, 0\n")
    difmapout.write("mapsize 1024," + str(pixsize) + "\n")
    difmapout.write("modelfit 50\n")
    difmapout.write("rmsflux = imstat(rms)\n")
    difmapout.write("restore\n")
    difmapout.write("pkflux = peak(flux,max)\n")
    difmapout.write("ilevs = pkflux/rmsflux\n")
    difmapout.write("lowlev = 300.0/ilevs\n")
    difmapout.write("loglevs lowlev\n")
    imagefile = rootdir + '/' + experiment + '/' + source + '.' + bands + \
                '.image.fits' + jmsuffix
    if os.path.exists(imagefile):
        os.system("rm -f " + imagefile)
    difmapout.write("wmap " + imagefile + "\n")
    difmapout.write("unshift\n")
    difmapout.write("wmod\n")
    difmapout.write("print rmsflux\n")
    difmapout.write("print imstat(bmin)\n")
    difmapout.write("print imstat(bmaj)\n")
    difmapout.write("print imstat(bpa)\n")

# JMFIT
def jmfit(pulsar, uvweightstr, experiment, jmsuffix):
    #Do for each IF/pol that all antennas were in on
    todo = ['rr.1', 'll.1', 'rr.2', 'll.2', 'rr.3', 'll.3', \
            'rr.4', 'll.4', 'combined']
    imagedata = AIPSImage('JUNK', 'IMG', 1, 1)
    outdata = AIPSImage('JUNK', 'IMG', 1, 2)
    messages = []
    for ifpol in todo:
        #First load the image data from difmap
        if imagedata.exists():
            imagedata.zap()
        if outdata.exists():
            outdata.zap()
        fitld = AIPSTask('fitld', version = aips_version_std)
        fitld.infile = rootdir + '/' + experiment + '/' + pulsar + '.' + \
                       ifpol + '.image.fits' + jmsuffix
        fitld.outdata = imagedata
        fitld.optype = ''
        fitld.ncount = 0
        fitld.dotable = 1
        fitld.douvcomp = -1
        fitld.doconcat = -1
        fitld()
    
        jmfit = AIPSTask('jmfit')
        jmfit.indata = imagedata
        jmfit.outdata = outdata
        jmfit.niter = 4000
        jmfit.ngauss = 1
        jmfit.fwidth[1][1] = 0
        jmfit.fwidth[1][2] = 0
        jmfit.fwidth[1][3] = 0
        jmfit.domax[1:] = [1]
        jmfit.dopos[1][1] = 1
        jmfit.dopos[1][2] = 1
        jmfit.dowidth[1][1] = 1
        jmfit.dowidth[1][2] = 1
        jmfit.dowidth[1][3] = 1
        jmfit.blc[1:] = [int(imagedata.header.crpix[0])-40, \
                         int(imagedata.header.crpix[1])-40]
        jmfit.trc[1:] = [int(imagedata.header.crpix[0])+40, \
                         int(imagedata.header.crpix[1])+40]
        jmfit.go()

        jmfitmessage = jmfit.message()
        msgindex = 0
        exciselinenos = []
        for i in range(len(jmfitmessage)-1):
            if jmfitmessage[len(jmfitmessage)-(i+1)][:5] != 'JMFIT':
                exciselinenos.append(len(jmfitmessage)-(i+1))
        for e in exciselinenos:
            jmfitmessage = jmfitmessage[:e] + jmfitmessage[e+1:]
        while jmfitmessage[msgindex].find('X-ref') == -1 and \
                  msgindex < len(jmfitmessage):
            msgindex = msgindex + 1
        if msgindex >= len(jmfitmessage):
            print "Did not find solution for " + pulsar + ", band " + ifpol
            continue
        centrerasplit = jmfitmessage[msgindex].split()
	centredecsplit = jmfitmessage[msgindex+1].split()
        msgindex = msgindex + 2
        while (jmfitmessage[msgindex].find('********* Solution from JMFIT') \
               == -1) and (msgindex < len(jmfitmessage)):
            msgindex = msgindex + 1
        if msgindex >= len(jmfitmessage):
            print "Did not find solution for " + pulsar + ", band " + ifpol
            continue
        sourcerasplit = jmfitmessage[msgindex+7].split()
        sourcedecsplit = jmfitmessage[msgindex+8].split()
        fluxsplit = jmfitmessage[msgindex+3].split()
        beammajsplit = jmfitmessage[msgindex+12].split()
        beamminsplit = jmfitmessage[msgindex+13].split()
        beampasplit = jmfitmessage[msgindex+14].split()
        try:
            beammaj = float(beammajsplit[4])
            beammin = float(beamminsplit[4])
        except ValueError:
            print "Bad value for one or more of major axis " +  \
                  beammajsplit[4] + ", minor axis " + beamminsplit[4] + \
                  ", both will be set to zero"
            beammaj = 0.0
            beammin = 0.0
        
        messages.append("Pulsar " + pulsar + ": Frequency " + ifpol[-1] + \
                        ", pol " + ifpol[:2].upper() + ":")
        messages.append("Centre RA:            " + centrerasplit[5] + ":" + \
                        centrerasplit[6] + ":" + centrerasplit[7])
        messages.append("Centre Dec:           " + centredecsplit[5] + ":" + \
                        centredecsplit[6] + ":" + centredecsplit[7])
        if fluxsplit[4] == '+/-':
            flux = float(fluxsplit[3][1:])*1000.0
            rms = float(fluxsplit[5])*1000.0
        else:
            flux = float(fluxsplit[4])*1000.0
            rms = float(fluxsplit[6])*1000.0
        messages.append("Flux (mJy):           " + str(flux))
        messages.append("S/N:                  " + str(flux/rms))
        centrerahours = float(centrerasplit[5]) + \
                        float(centrerasplit[6])/60.0 + \
                        float(centrerasplit[7])/3600.0
        centredecdegs = float(centredecsplit[5]) + \
                        float(centredecsplit[6])/60.0 + \
                        float(centredecsplit[7])/3600.0
        srcrahours = float(sourcerasplit[2]) + \
                     float(sourcerasplit[3])/60.0 + \
                     float(sourcerasplit[4])/3600.0
        srcdecdegs = float(sourcedecsplit[2]) + \
                     float(sourcedecsplit[3])/60.0 + \
                     float(sourcedecsplit[4])/3600.0
        raoff = (srcrahours-centrerahours)*15*60*60*1000*\
                math.cos(centredecdegs*math.pi/180.0)
        decoff = (srcdecdegs-centredecdegs)*60*60*1000
        messages.append("RA offset (mas):      " + str(raoff))
        messages.append("Dec offset (mas):     " + str(decoff))
        messages.append("Actual RA:            " + sourcerasplit[2] + ":" + \
                        sourcerasplit[3] + ":" + sourcerasplit[4])
        messages.append("Actual Dec:           " + sourcedecsplit[2] + ":" + \
                        sourcedecsplit[3] + ":" + sourcedecsplit[4])
        messages.append("Beam:                 " + str(beammin*1000) + "x" + \
                         str(beammaj*1000) + " at " + beampasplit[4] + \
                        " degrees")
        raerr = float(sourcerasplit[6])*1000
        decerr = float(sourcedecsplit[6])*1000
        messages.append("Est. RA error (mas):  " + \
                        str(raerr*math.cos(centredecdegs*math.pi/180.0)*15.0))
        messages.append("Est. RA error (hms):  " + str(raerr))
        messages.append("Est. Dec error (mas): " + str(decerr))
    return messages

for pulsar in targets:
    (ra, dec) = get_ra_dec(pulsar, uvprefix, uvsuffix)
    print "Ra is " + str(ra) + ", dec is " + str(dec)
    difmap_mappsr(pulsar, isll, ra, dec, uvweightstr, experiment, \
                  difmappath, uvprefix, uvsuffix, jmsuffix, doaddnoise)
    messages = jmfit(pulsar, uvweightstr, experiment, jmsuffix)
    jmfitstatsfilename = rootdir + '/' + experiment + '/' + \
                         pulsar + ".jmfitstats" + jmsuffix
    jmfitstatsout = open(jmfitstatsfilename, "w")
    for line in messages:
        jmfitstatsout.write(line + "\n")
    jmfitstatsout.close()
    if doaddnoise:
        difmap_mappsr(pulsar, isll, ra, dec, uvweightstr, experiment, \
                      difmappath2, "noise", ".uvf", jmsuffix+".nowt", False)
        messages = jmfit(pulsar, uvweightstr, experiment, jmsuffix+".nowt")
        jmfitstatsfilename = rootdir + '/' + experiment + '/' + \
                             pulsar + ".jmfitstats" + jmsuffix + ".nowt"
        jmfitstatsout = open(jmfitstatsfilename, "w")
        for line in messages:
            jmfitstatsout.write(line + "\n")
        jmfitstatsout.close()