import glob, os, re

def summary_start(weighting, dec):
    outfile = "almabeams_%s_%s.summary" % (weighting, dec)
    try:
        ofile = open(outfile, 'w')
    except:
        print "Error opening", ofile, "for writing."
        return ofile

    ofile.write("%s weighted 672 GHz 4h obs., dec = %s deg\n" % (weighting,
                                                                 dec))
    ofile.write('Cfg.    Min    Max    RMS   Maj (mas)  Min (mas)  Rel. noise\n')
    return ofile


def write_beam_stats(bim, ofile, relsens, beam, confnum):
    ofile.write("O-%02d" % confnum)
    ia.open(bim)
    bstats = ia.statistics(robust=1,verbose=false)
    ia.close()
    ofile.write("  %6.3f  %5.3f  %6.4f" % (bstats['min'][0],
                                           bstats['max'][0],
                                           bstats['medabsdevmed'][0]))
    ofile.write("  %7.3f   %7.3f" % (1000.0 * beam['bmaj'].value,
                                     1000.0 * beam['bmin'].value))
    ofile.write("  %5.3f\n" % relsens)


def do_other_weightings(dec, avgsepsf="../configs.avgsep"):
    """
    Make robust and uniform weighted dirty beams for measurement sets have
    already been made at declination dec.
    """
    # read avgseps.
    try:
        avsfile = open(avgsepsf, 'r')
        avgseps = map(float, avsfile.readlines())
        avsfile.close()
    except:
        print "Could not read", avgsepsf
        return

    rsumm = summary_start('robst1', dec)
    usumm = summary_start('unifrm', dec)

    direcshun = "J2000 18h00m00.03s %sd0m0.0s" % dec
    for confnum in range(1, 29):
        confname = "out%02d" % confnum
        msfn = glob.glob("*%s_4h*_nat*_nn.ms" % confname)[0]
        im.open(msfn)

        pixelsize = 0.1 * avgseps[0] / avgseps[confnum - 1]
        cell = "%farcsec" % pixelsize
        im.defineimage(nx=256, ny=256, cellx=cell, phasecenter=direcshun)

        im.weight(type='natural')
        nsens = im.sensitivity()['pointsource'].value

        rdbim = re.sub('nat[^_]*', 'robst1', msfn)
        rdbim = re.sub('ms$', 'psf', rdbim)
        im.weight(type='briggs',rmode='norm',robust=1.0)
        rsens = im.sensitivity()['pointsource'].value
        if os.path.isfile(rdbim) is False:
            im.makeimage(type='psf',image=rdbim)
        rbeam=im.fitpsf(psf=rdbim)
        
        im.weight(type='uniform')
        usens = im.sensitivity()['pointsource'].value
        udbim = re.sub('robst1', 'unifrm', rdbim)
        if os.path.isfile(udbim) is False:
            im.makeimage(type='psf',image=udbim)
        ubeam=im.fitpsf(psf=udbim)
        im.done()

        write_beam_stats(rdbim, rsumm, rsens / nsens, rbeam, confnum)
        write_beam_stats(udbim, usumm, usens / nsens, ubeam, confnum)

    rsumm.close()
    usumm.close()


def writebeams():
    """
    Export beam images to FITS.
    """
    dec = -23
#    for dec in (20, -85):
    for ttime in ('30s', '14400s'):
        for confnum in range(1, 29):
            confname = "out%02d" % confnum
            print "confname:", confname
            casaim = "testalmaconfig_%s_%s_%s_briggs_nn.psf" % (confname,
                                                                ttime, dec)
            outfile = "%s_briggs_%s_%s_psf.fits" % (confname, ttime, dec)
            exportfits(imagename = casaim, fitsimage = outfile, async=True)


def summarize_beams(outfileroot='almabeams_briggs'):
    """
    Write a summary of the ALMA beams into outfileroot_dec.summary.
    """
    dec = -23
#    for dec in (20, -85):
    for ttime in ('30s', '14400s'):
        outfile = outfileroot + "_%s_%s.summary" % (ttime, dec)
        try:
            ofile = open(outfile, 'w')
        except:
            print "Error opening", ofile, "for writing."
            return

        ofile.write("Briggs weighted 672 GHz %s obs., dec = %s deg\n" % (ttime,
                                                                         dec))
        ofile.write('Cfg.    Min    Max    RMS   Maj (mas)  Min (mas)\n')
        for confnum in range(1, 29):
            confname = "out%02d" % confnum
            ofile.write("O-%02d" % confnum)
            casaim = "testalmaconfig_%s_%s_%s_briggs_nn.psf" % (confname,
                                                                ttime, dec)
            ia.open(casaim)
            bstats = ia.statistics(robust=1,verbose=false)
            ia.close()
            ofile.write("  %6.3f  %5.3f  %6.4f" % (bstats['min'][0],
                                                   bstats['max'][0],
                                                   bstats['medabsdevmed'][0]))
            casaim = "testalmaconfig_%s_%s_%s_briggs_nn.restored" % (confname,
                                                                     ttime, dec)
            ia.open(casaim)
            beam = ia.restoringbeam()
            ia.close()
            ofile.write("  %7.3f   %7.3f" % (1000.0 * beam['major']['value'],
                                             1000.0 * beam['minor']['value']))
            ofile.write('\n')
        ofile.close()