import numpy
from math import sqrt

# Originally written by Antonio Hales.  Spiffied up and translated to python
# with epydoc documenation by Rob Reid.
def suggestalmaconfig(res, frequency, declination):
    """
    Returns a suggested ALMA configuration for a desired resolution and
    frequency.  It assumes a 4h observation, but the observation length should
    not matter.

    @param res: desired FWHM naturally weighted dirty beam major axis in \".
    @type res: float or int
    @param frequency: Center frequency to use in GHz.
    @type frequency: float or int
    @param declination: declination to use in degrees.
    @type declination: float or int
    """

    # Declinations with precalculated beam sizes.
    # Must match summary filenames (no .0s).
    declins = [-85, -46, -23, 0, 20]

    southdecind = 0
    while southdecind < 3 and declination > declins[southdecind + 1]:
        southdecind += 1
        
    # Calculate polated resolutions.
    sdec = declins[southdecind]
    ndec = declins[southdecind + 1]
    sfwhms = beamsizes(sdec)
    nfwhms = beamsizes(ndec)
    pfwhms = ((ndec - declination) * sfwhms +
              (declination - sdec) * nfwhms) / (ndec - sdec)

    res *= 672.0 / frequency
    mindiff = 3.0 * pfwhms[0]
    bestcfg = 0
    for i in range(0, len(pfwhms)):
        diff = abs(res - pfwhms[i])
        if diff < mindiff:
            mindiff = diff
            bestcfg = i
    
    return bestcfg + 1


def beamsizes(declin):
    """
    Reads almabeams<declin>.summary and returns an array of beam sizes in
    arcseconds, going from out01 to out28 (= Y8).  The beam size is the FWHM
    square root of the minor axis x the major axis.

    @param declin: declination to use in degrees.
    @type declin: int
    """
    sfn = "almabeams_%s.summary" % declin
    if os.path.isfile(sfn) is False:
        sfn = os.popen('locate %s' % sfn).readline().rstrip()
    sfile = open(sfn, 'r')
    sfn.close()
    fwhms = []
    for line in sfile:
        cols = line.split()
        try:
            fwhms.append(0.001 * sqrt(float(cols[4]) * float(cols[5])))
        except:
            pass
    sfile.close()
    return numpy.array(fwhms)