"""Function objects for primary beams of circular apertures."""

import scipy
from scipy.special import j1

class airypat:
    def __init__(self, a, freq):
        """
        Returns an Airy *pattern* function object at frequency freq GHz
        for a circular aperture with radius a meters, which takes as its
        argument the sine of the angle from the optical axis.
        """
        self.a    = a
        self.freq = freq
        self.rfac = 20.959 * a * freq    # 20.959 ~= 2pi GHz/c m^-1

    # Handle r when it is a scalar.
    def _scal_handler(self, r):
        if r < 0.0:               # %s more general than %f.
            raise ArithmeticError, "%s < 0 cannot be a radius." % r
        
        if r < 0.001:               # Handle 0 / 0 limit
            ro2sq = 0.25 * r**2
            # Good to +-6.78e-31
            ans = 1.0 + ro2sq * (-0.5 + ro2sq * (1.0 / 12.0 - ro2sq / 144.0))
        else:
            ans = 2.0 * j1(r) / r
        return ans
        
    # Handle r when it is a vector.
    def _vect_handler(self, r):
        smallr_inds = scipy.where(r < 0.001)
        bigr_inds   = scipy.where(r > 0.001)

        smallrs = r[smallr_inds]
        bigrs   = r[bigr_inds]
        
        ans = scipy.zeros(r.shape)               # Initialize ans as an array.
        ans[bigr_inds] = 2.0 * j1(bigrs) / bigrs
        
        # Handle 0 / 0 limit
        ro2sq = 0.25 * smallrs**2
        # Good to +-6.78e-31
        ans[smallr_inds] = 1.0 + ro2sq * (-0.5 + ro2sq * (1.0 / 12.0 - ro2sq / 144.0))

        return ans
        
    def __call__(self, salfa):
        """
        Returns the normalized* intensity value at salfa = sin(alfa) for a
        circular aperture.  The aperture radius and frequency should be set
        when the airypat instance is created.

        * i.e. airypat(salfa = 0) = 1.
        """
        r = self.rfac * salfa
        if isinstance(r, (float, int)):
            return self._scal_handler(r)
        else:                           # Assume it's an array.
            return self._vect_handler(r)

class pb_of_circaper:
    def __init__(self, a, b, freq):
        """
        Returns a function object for the primary beam at frequency freq GHz
        for a circular aperture with radius a meters and radius b meters of
        blockage in the center.  The argument of the returned function is
        arcseconds in the sin projection!
        """
        self.a    = a
        self.b    = b
        self.freq = freq

        # A sneaky way to convert arcseconds to radians.
        freq *= 4.84813681109536e-06

        unbl = airypat(a, freq)
        
        if b > 0.0:
            if b >= a:
                raise ArithmeticError, "There is no unblocked aperture!"
            blked = airypat(b, freq)
            self.func = lambda r: ((a * unbl(r) - b * blked(r)) / (a - b))**2
        else:
            self.func = lambda r: unbl(r)**2

    def __call__(self, r):
        """
        Returns the normalized* intensity value at r arcseconds for a
        possibly blocked circular aperture.  The aperture radii and frequency should be set
        when the pb_of_blockeddisk instance is created.

        * i.e. pb_of_blockeddisk(r = 0) = 1.
        """
        return self.func(r)