"""
psf.py 

psf models
"""

import os

import numpy as np 

from numpy.polynomial.polynomial import polyvander2d, polyval2d
from numpy.polynomial.legendre import leggauss
from scipy.interpolate import interp1d, RectBivariateSpline

from ..util.io import NTuple
import scipy.interpolate
from ..lib.fitparameters import FitParameters
from ..lib.instruments import InstrumentModel, MagSys
# TODO: move load_ascii into util or io
from ..lib.spectrum import Spectrum, load_ascii

import logging
logging.basicConfig(format="%(asctime)-15s %(levelname)s %(message)s",
                    level=logging.INFO)
    

def integ_in_pixels(x, y, f, quadrature_order=4):
    """Integrate a 2D function in pixels
    
    Given a function f (typically, a psf profile), and a series of
    positions (x,y), evaluate the integral of the function on the
    pixels corresponding to each of the positions (x,y) specified in
    argument.

    Args
       x (ndarray of floats): x-positions, in pixels.
       y (ndarray of floats): y-positions, in pixels.
       f (callable): function to integrate
       quadrature_order (int, optional): order of the gaussian quadrature to use.

    Returns:
       ndarray of floats: integrated PSF value.

    .. note: 
       the function is written to preserve the shape of x and y
    """
    # put this integration code in cache
    p,w = leggauss(quadrature_order)
    p *= 0.5 ; w *= 0.5
    gx,gy = np.meshgrid(p,p)
    wx,wy = np.meshgrid(w,w)
    ww = (wx*wy).ravel()
    #    xc = (np.floor(x)+0.5).reshape(x.shape + (1,)) + p
    #    yc = (np.floor(y)+0.5).reshape(y.shape + (1,)) + p
    xc = x.reshape(x.shape + (1,1,)) + gx
    yc = y.reshape(y.shape + (1,1,)) + gy
    v = f(xc,yc)
    v = v.reshape(v.shape[0], v.shape[1], -1)
    return np.dot(v, ww)


class GaussianPsf(object):
    """
    A class for a 2D gaussian PSF

    :math:
       PSF(x,y) = F \times \frac{|W|^{1/2}}{2\pi} \exp(-\frac{1}{2} (X-X_0)^T W (X-X_0))
    """    
    def __init__(self):
        self.pars = self.init_pars()        

    def init_pars(self):
        """Default parameters for a gaussian PSF.

        By default, the PSF is centered and normalized to 1.

        Args:
            seeing (float): seeing value
            fwhm (float, optional): FWHM
        """
        p = FitParameters(['flux', 'x0', 'y0', 'wxx', 'wyy', 'wxy'])
        p['flux'], p['x0'], p['y0'] = 1., 0., 0.
        p['wxx'], p['wyy'], p['wxy'] = 0.25, 0.25, 0.
        return p
    
    @property
    def fwhm(self):
        wxx, wyy = self.pars['wxx'].full[0], self.pars['wyy'].full[0]
        ww = np.sqrt(wxx*wyy)
        return 2. * np.sqrt(2. * np.log(2.) / ww)

    @fwhm.setter
    def fwhm(self, v):
        old_fwhm = self.fwhm
        scale = (old_fwhm / v)**2
        self.pars['wxx'].full *= scale
        self.pars['wyy'].full *= scale
        self.pars['wxy'].full *= scale

    @property
    def sigma(self):
        wxx, wyy = self.pars['wxx'].full[0], self.pars['wyy'].full[0]
        return 1. / np.power(wxx*wyy, 0.25)

    @sigma.setter
    def sigma(self, v):
        old_sigma = self.sigma
        scale = (old_sigma / v) ** 2
        self.pars['wxx'].full *= scale
        self.pars['wyy'].full *= scale
        self.pars['wxy'].full *= scale
        
    def profile(self, x, y, fp=None):
        """PSF profile at positions (x,y)

        Evaluate and return the PSF profile at the specified
        positions. The function is written so that the shape of 
        the input arrays (x,y) is preserved.
        
        Args:
            x (ndarray of floats): x-positions (pixels)
            y (ndarray of floats): y-positions (pixels)
            fp (FitParameters, optional): parameters of the gaussian

        Returns:
            ndarray of floats: the values of the (non integrated) profile
            
        .. note:
            on a Intel Core i5-2540M CPU @ 2.60GHz, about 100 micro-sec 
            to evaluate on ~ 1800 points.
        """
        if fp is None:
            fp = self.pars
        flux, x0, y0, wxx, wyy, wxy = fp.full
        det = wxx*wyy-wxy**2
        norm = np.sqrt(det) / (2.*np.pi)
        xx, yy = x-x0, y-y0
        r = -0.5 * (xx*xx*wxx + yy*yy*wyy + 2.*wxy*xx*yy)
        return flux * norm * np.exp(r)
    
    def __call__(self, x, y, fp=None):
        """psf value, integrated in the pixels.
        
        evaluate the psf value, integrated in each of the pixels 
        where x,y belong.
        
        Args:
            x (ndarray of floats): x-positions (pixels)
            y (ndarray of floats): y-positions (pixels)
            fp (FitParameters, optional): PSF parameters
               if not present, the PSF uses its own internal 
               parameter vector. If is safer to use a PSF
               parameter generated by self.init_pars()

        Returns:
            array of floats: psf values.

        .. note: on a Intel Core i5-2540M CPU @ 2.60GHz, about 750
            micro-sec to evaluate on ~ 1800 points. Almost half of it
            is the call to leggauss in integ_in_pixels
        """
        if fp is not None:
            self.pars.full[:] = fp.full[:]
        if type(x) is float or type(x) is int:
            x = np.array([x])
        if type(y) is float or type(y) is int:
            y = np.array([y])
        return integ_in_pixels(x, y, self.profile)
    
    def pixels(self, rad=5., integrate=True, fp=None, dx=None):
        """array of pixels, containing the PSF values.
        
        Return an array of pixels, with the PSF values. By default, 
        these values are integrated in the pixels. 
        
        Args:
            r (float): radius, in sigma units
            integrate (bool, optional): if true, integrate the PSF profile.
            fp (FitParameters, optional): PSF fit parameters

        Returns:
            ndarray of floats: contains the PSF values
        """
        if fp is None:
            fp = self.pars
        if dx is None:
            dx = int(rad * self.fwhm / 2.)
        x = np.arange(-dx, dx+1., 1.)
        x,y = np.meshgrid(x,x, sparse=True)
        if integrate:
            pix = self(x,y, fp=fp)
        else:
            pix = self.profile(x,y, fp=fp)
        return pix / pix.sum()
        

class MoffatPsf(object):
    """
    A class for a MoffatPSF
    
    :math:
       PSF(X) = \frac{F}{(1 + (X-X_0)^T W (X-X_0))^\alpha}
    """
    def __init__(self, seeing=2., fwhm=None, expn=2.5):
        self.pars = self.init_pars(expn)
        self.fwhm = 4.7096
        
    def init_pars(self, expn):
        """return a initialized parameter vector for the psf. 
        
        Args:
            seeing (float): sigma, in pixels 
            fwhm (float, optional): fwhm, in pixels
            expn (float): Moffat exponent

        Returns:
            FitsParameters: a vector of type FitParameters, 
                initialized so that the PSF has the requested seeing or fwhm
        """
        p = FitParameters(['flux', 'x0', 'y0', 'wxx', 'wyy', 'wxy', 'expn'])
        p['flux'], p['expn'] = 1., expn
        p['wxx'], p['wyy'], p['wxy'] = 0.05761906, 0.05761906, 0.
        return p

    @property
    def fwhm(self):
        expn = self.pars['expn'].full[0]
        wxx, wyy = self.pars['wxx'].full[0], self.pars['wyy'].full[0]
        ww = np.sqrt(wxx*wyy)
        return 2. * np.sqrt((np.power(2., 1/expn) - 1.) / ww)
    
    @fwhm.setter 
    def fwhm(self, v):
        old_fwhm = self.fwhm
        scale = (old_fwhm / v) ** 2
        self.pars['wxx'].full *= scale
        self.pars['wyy'].full *= scale
        self.pars['wxy'].full *= scale        
        
    def profile(self, x, y, fp=None):
        """
        return the (non-integrated) PSF profile for the specified x,y positions
        
        Args:
            x (ndarray of floats): x-positions in pixels
            y (ndarray of floats): y-positions in pixels            
            fp (FitParameters, optional): parameter vector

        Returns:
            the psf profile values at positions x,y

        .. note:
            the shape of x,y is preserved
        """
        if fp is None:
            fp = self.pars
        flux, x0, y0, wxx, wyy, wxy, expn = fp.full
        det = wxx*wyy-wxy**2        
        norm = np.sqrt(np.fabs(det)) * (expn-1.) / np.pi
        xx,yy = x-x0, y-y0
        fact = 1. / (1. + xx*xx*wxx + yy*yy*wyy + 2*xx*yy*wxy)
        return norm * flux * np.power(fact, expn)

    def __call__(self, x, y, fp=None):
        if fp is not None:
            self.pars.full[:] = fp.full[:]
        if type(x) is float or type(x) is int:
            x = np.array([x])
        if type(y) is float or type(y) is int:
            y = np.array([y])
        return integ_in_pixels(x, y, self.profile)

    def pixels(self, rad=5., integrate=True, fp=None, dx=None):
        if fp is None:
            fp = self.pars
        if dx is None:
            dx = int(rad * self.fwhm / 2.)
        x = np.arange(-dx, dx+1., 1.)
        x,y = np.meshgrid(x,x, sparse=True)
        if integrate:
            pix = self(x,y, fp=fp)
        else:
            pix = self.profile(x,y, fp=fp)
        return pix / pix.sum()

def new_psf(name):
    import re
    if name == 'GAUSSIAN':
        return GaussianPsf()
    expr = re.search('MOFFAT(\d+)', name)
    if expr is not None:
        try:
            beta = float(expr.group(1)) * 0.1
            assert beta > 0.
            return MoffatPsf(expn=beta)
        except:
            pass
    return None


class PsfFluxFisher(object):
    """Efficient computation of the variance of a PSF flux estimator.
    
    The variance of the flux is given by: 
    :math:
        V(f) = \left(\sum_i\frac{\psi_i^2}{f \psi_i + f_{sky|i} + \sigma^2_{RON}}\right)^{-1}

    The class precomputes the values of function:
    :math:
        W(x,\sigma) = \sum_i \frac{\psi_i^2}{\psi_i + x}

    over a user-defined range in fwhm, and a range
    :math:`[10^{-6},10^6]` in x. 

    By interpolating (linearly) over the precomputed values one can
    compute very efficiently the variance of the measured PSF flux,
    given the flux (in e-/s), the sky value (e-/s/pixel) and the
    seeing (fwhm, pixels):
    :math:
       V(f) = \frac{f}{W(\frac{f_{sky} + \sigma^2_{RON}}{f}, fwhm)}

    .. note: the class does not know anything about pixel size,
             arcseconds, gain, exposure times and the like. 
             
             Fluxes and sky fluxes are in electrons, seeing is fwhm in
             pixels.  The conversion from physical quantities is
             performed elsewhere (in the instrument class).
    """
    def __init__(self, psf_model=None):
        if psf_model is not None:
            self.psf_model = psf_model
        else:
            self.psf_model = GaussianPsf()

    def precompute(self, x_range=(-6., 6.), fwhm_range=(0.5, 15.), full=False):
        self._precompute_W(x_range=x_range, fwhm_range=fwhm_range)
        if full:
            self._precompute_U()
                    
    def _precompute_W(self, x_range=(-6, 6.), fwhm_range=(0.5, 15.)):
        """precompute the auxiliary function W(x,fwhm)
        
        Loop over the x- and fwhm-ranges, and precompute the W(x,fwhm)
        values over that grid.

        In fact, what is pre-computed and stored is: 
        :math:
           \log_{10} W(\log_10(x), fwhm)

        .. note: This function is somewhat slow (typically, half a second), 
                 but needs to be called only once in the life of the object
                 (unless one decides to change the PSF model, in which case 
                 it has to be called again)
        """
        logging.info('ETC: precomputing W() function [takes a few secs]')
        self.fwhm_range = fwhm_range
        self.x_range = x_range
        fwhm = np.arange(fwhm_range[0], fwhm_range[1], 0.1)
        x = np.logspace(x_range[0], x_range[1], 100)
        x = x.reshape(x.shape + (1,))
        ww, invnea = [], []
        # we loop over the fwhm values first, 
        # to minimize the recomputations of the pixel grids.
        for s in fwhm:
            self.psf_model.fwhm = s
            pix = self.psf_model.pixels(integrate=True).reshape(1, -1)
            val = np.sum(pix**2 / (pix + x), axis=1)
            ww.append(np.log10(val).ravel())
            invnea.append(np.sum(pix**2))
        self.x = x, self.fwhm = x, fwhm
        self.ww = np.vstack(ww)
        self._W = RectBivariateSpline(np.log10(x), fwhm, self.ww.T)
        # need to interpolate between the NEA
        self.invnea = np.asarray(invnea)
        self._invnea = interp1d(fwhm, self.invnea)        

    def _precompute_U(self):
        """
        to be called after precompute_W
        """
        if not hasattr(self, '_W'):
            self._precompute_W()
        logging.info('ETC: precomputing U() function [takes a few secs]')
        from scipy.optimize import brenth
        def f(lx, seeing, lp):
            return self._W.ev(lx,seeing) - lx - lp
        fwhm = np.arange(self.fwhm_range[0], self.fwhm_range[1], 0.1)
        lp = np.linspace(-5., 5., 200)
        gp,gs = np.meshgrid(lp,fwhm)
        data = []
        lx_min, lx_max = -6., 6.
        for lp_,s_ in zip(gp.ravel(),gs.ravel()):
            lx = brenth(f, lx_min, lx_max, args=(s_,lp_))
            data.append(lx)
        self.u_data = data = np.asarray(data)
        self._U = RectBivariateSpline(lp, fwhm, data.reshape(gs.shape).T)
        
    def W(self, x, fwhm):
        """evaluate W for the values of x and fwhm passed as args.

        Evaluate 

        Args:
            x (ndarray of floats): x-values
            fwhm (ndarray of floats): y-values

        Returns:
            ndarray of floats: contains the values of W 
        
        Raises: 
            AssertionError if precompute has not been called.
        """
        assert len(x) == len(fwhm)
        if not hasattr(self, '_W'):
            self._precompute_W()
        x_min, x_max = 10**self.x_range[0], 10**self.x_range[1]
        s_min, s_max = self.fwhm_range
        lx = np.log10(x)
        r = 10**self._W.ev(lx,fwhm)
        # if x outside interp bounds, we have asymptotic values
        r[x<x_min] = 1.
        idx = x>x_max
        r[idx] = self._invnea(fwhm[idx]) / (x[idx])
        # if fwhm outside interp bounds, return very low weight 
        # why not zero ? 
        r[(fwhm<s_min) | (fwhm>s_max)] = 0. # 1.E-9
        return r        
    
    def flux_variance(self, flux, sky, fwhm, ron=5.):
        """compute the PSF flux variance

        Compute the variance of the PSF flux estimator, given values 
        for the flux, sky, seeing (in pixels) and readout noise.
        
        Args:
            flux (ndarray of floats): object fluxes (in electrons)
            sky (ndarray of floats): sky background value (in electrons)
            fwhm (ndarray of floats): fwhm of the PSF (pixels)
            ron (float): readout noise (electrons)
            
        Returns:
            ndarray of floats: flux variance
        """
        fwhm = np.asarray(fwhm)
        flux = np.asarray(flux)
        sky = np.asarray(sky)
        x = (sky + ron**2) / flux
        w = self.W(x, fwhm)
        return flux / w
    
    def noise_equivalent_area(self, fwhm):
        if not hasattr(self, '_invnea'):
            self._precompute_W()
        return 1. / self._invnea(fwhm)

    def signal_to_noise(self, flux, sky, fwhm, ron=5.):
        """return the signal-to-noise ratio

        """
        fwhm = np.asarray(fwhm)
        flux = np.asarray(flux)
        sky = np.asarray(sky)
        x = (sky + ron**2) / flux 
        return np.sqrt(flux * self.W(x, fwhm))

    def flux_lim(self, sky, fwhm, ron=5., snr=5.):
        """return the limiting flux that correspond to the snr specified in argument

        Args:
            sky (ndarray of floats): sky level, in electrons
            fwhm (ndarray of floats): PSF fwhm, in pixels

        Returns:
            ndarray of floats: limiting flux
        """
        if not hasattr(self, '_U'):
            self._precompute_U()
        log_snr_lim = np.log10(snr**2 / (sky+ron**2))
        log_x_lim = self._U.ev(log_snr_lim, fwhm)
        return (sky + ron**2) * 10**(-log_x_lim)
        
    def plot_variance(self, flux_range=(0.1, 1E9), sky=5000., fwhm=4., ron=5., npts=200, 
                      slow=False, fmt='b-'):
        import pylab as pl 
        f_min, f_max = list(map(np.log10, flux_range))
        flux_ = np.logspace(f_min, f_max, npts)
        sky_ = np.ones(npts) * sky
        fwhm_ = np.ones(npts) * fwhm
        if slow is False:
            V = self.flux_variance(flux_, sky_, fwhm_, ron)
        else:
            V = self.flux_variance_slow(flux_, sky_, fwhm_, ron)
        pl.loglog(flux_, V, fmt)
        pl.xlabel('Flux [e-]')
        pl.ylabel('V(Flux)')
        pl.title('PSF flux variance [sky=%6.1f e-, fwhm=%3.1f ron=%4.2f]' % (sky, fwhm, ron))

    def plot_snr(self, flux_range=(0.1, 1.E9), sky=5000., fwhm=4., ron=5., npts=200,
                 slow=True, fmt='b-'):
        import pylab as pl
        f_min, f_max = list(map(np.log10, flux_range))
        flux_ = np.logspace(f_min, f_max, npts)
        sky_ = np.ones(npts) * sky
        fwhm_ = np.ones(npts) * fwhm
        if slow is False:
            snr = self.signal_to_noise(flux_, sky_, fwhm_, ron)
        else:
            sig = np.sqrt(self.flux_variance_slow(flux_, sky_, fwhm_, ron))
            snr = flux_ / sig            
        pl.loglog(flux_, snr, fmt)
        pl.xlabel('Flux [e-]')
        pl.ylabel('SNR(Flux)')        
        pl.title('PSF flux SNR [sky=%6.1f e-, fwhm=%3.1f ron=%4.2f]' % (sky, fwhm, ron))
    
    def plot_fluxlim(self, snr=5., sky_range=(1., 1.E5), fwhm=4., ron=5., npts=200, fmt='b-'):
        import pylab as pl
        sky_min, sky_max = list(map(np.log10, sky_range))
        sky_ = np.logspace(sky_min, sky_max, npts)
        fwhm_ = np.ones(npts) * fwhm
        flux_lim = self.flux_lim(sky_, fwhm_, ron=ron, snr=snr)
        pl.plot(sky_, flux_lim, fmt)
        pl.xlabel('sky level [e-/pixel]')
        pl.ylabel('Flux lim [SNR=%5.1f]' % snr)        
        pl.title('PSF Flux lim [fwhm=%3.1f ron=%4.2f]' % (fwhm, ron))
        
    def flux_variance_slow(self, flux, sky, seeing, ron=5.):
        """compute the PSF flux variance (slow version)

        Same as above, but does not interpolate through precomputed
        values. Use it for debugging purposes only.

        .. note: on a Intel Core i5-2540M CPU @ 2.60GHz, about 
                 2500 times slower than function above. 
        """
        r = []
        for s, f, sk in zip(seeing, flux, sky):
            self.psf_model.fwhm = s
            pix = self.psf_model.pixels(integrate=True).ravel()
            v = np.sum(pix**2 / (f * pix + sk + ron**2))
            r.append(1./v)
        return np.array(r)

    def W_slow(self, fwhm, x):
        r = []
        for s,xx in zip(fwhm, x):
            self.psf_model.fwhm = s
            pix = self.psf_model.pixels(integrate=True).ravel()
            v = np.sum(pix**2 / (pix + xx))
            r.append(1. / v)
        return np.array(r)


class ImagingInstrument(object):
    """Exposure Time Calculator (ETC) for a specific imager.
    
    An instance of ImagingInstrument stores the characteristics of an
    instrument (passbands, pixel size, readout noise) and site
    (seeing, average sky brightness etc.)  Compute flux variance
    and limiting magnitudes for given observing conditions.
    
    An imaging instrument is probably the most 
    
    """
    def __init__(self, instrument_name, **kwargs):
        self.instrument = InstrumentModel(instrument_name)
        cards = self.instrument.cards
        # general cards
        self.ron = self._kw_or_card('ro_noise', kwargs, cards)
        self.dark_current = self._kw_or_card('dark_current', kwargs, cards)
        self.pixel_size = self._kw_or_card('pix_size', kwargs, cards)
        self.pixel_area = self.pixel_size ** 2
        
        # default IQ
        self.iq_model = self._load_default_iq_model()
        
        # default sky background model
        self.sky_background = self._load_default_sky_background()
        
        # psf and etc 
        psf_type = self._kw_or_card('PSF_MODEL', kwargs, cards, default='GAUSSIAN')
        self.psf = new_psf(psf_type)
        self.etc = PsfFluxFisher(self.psf)        
        #        self.etc.precompute(full=1)

        # key data per band 
        self.magsys = MagSys(kwargs.get('magsys', 'AB'))
        self.data = self._compute()
        self.zp = dict(list(zip(self.data['band'], self.data['zp'])))

    def _compute(self):
        fw = self.filter_wheel = self.instrument.filters
        band = list(fw.transmissions.keys())
        tr = [self.instrument.EffectiveFilterByBand(k) for k in band]
        zp = [self.magsys.ZeroPoint(t) for t in tr]
        #        self.zp = dict(zip(band, zp))
        leff = [t.mean(t.x_min, t.x_max) for t in tr]
        iq = [float(self.iq_model(l)) for l in leff]
        msky = [self.magsys.mag(t, self.sky_background)[0][0] for t in tr]
        fsky = [self.sky_background.IntegFlux(t) * self.pixel_area for t in tr]
        #  nea = [self.etc.noise_equivalent_area(s/self.pixel_size) for s in iq]
        #  exptime = [exptime] * len(band)        
        #  m5 = [float(self.mag_lim(t,m,s,b,snr=5.)) for t,m,s,b in zip(exptime,msky,iq,band)]
        full_band_name = [self.instrument.name + '::' + k for k in band]
        ret = np.rec.fromarrays((full_band_name, zp, leff, iq, msky, fsky, tr), 
                                dtype=[('band', '|U20')] + \
                                    [(n,'f8') for n in ('zp', 'leff', 'iq', 'mag_sky', 'flux_sky')] + \
                                    [('tr', 'O')])
        ret.sort(order='leff')
        return ret

    def _kw_or_card(self, kw, kwargs, cards, default=None, raise_exc=False):
        """utility function to interpret the datacards and kw arguments
        """
        r = kwargs.get(kw, None)
        if r is None:
            r = cards.get(kw.upper(), default)
            if r is None and raise_exc:
                raise ValueError('no value specified for: %s' % kw)
        try:
            r = int(r)
        except:
            try:
                return float(r)
            except:
                return r
        return r

    def _load_default_iq_model(self):
        if 'IQ_MODEL' not in self.instrument.cards:
            self.iq = scipy.interpolate.interp1d((3000., 12000.), (1., 1.))
        r = self.instrument.cards['IQ_MODEL'].split()
        if r[0] == 'IQTABLE':
            nm = self.instrument.path + os.sep + r[1]
            d = []
            with open(nm) as f:
                for line in f.readlines():
                    if line[0] == '#': 
                        continue
                    d.append(list(map(float, line.split())))
            iq = np.rec.fromrecords(d, names=['wl', 'iq'])
            return scipy.interpolate.interp1d(iq['wl'], iq['iq'])
        else:
            raise NotImplementedError('model: %s' % r[0])        

    def _load_default_sky_background(self):
        cards = self.instrument.cards
        assert 'SKY_BACKGROUND' in cards
        if cards['SKY_BACKGROUND'] == 'MAUNAKEA_GREY':
            return load_ascii(self.instrument.path + os.sep + '../../maunakea_sky' + os.sep + 'skyspectrum.dat')
        elif cards['SKY_BACKGROUND'] == 'ZODIACAL_LIGHT':
            return load_ascii(self.instrument.path + os.sep + '../../zodiacal_light' + os.sep + 'zodiacal_light.dat')
        else:
            raise ValueError('unable to load sky background')
        
    def _broadcast(self, *args):
        a = tuple([np.asarray(x).reshape(-1) for x in args])
        return np.broadcast_arrays(*a)        

    def precompute(self, full=False):
        self.etc.precompute(full=full)

    def zero_points(self, band):
        return np.asarray([self.zp[b] for b in band])
    
    def flux_to_mag(self, flux, band, zp=None):
        if zp is None:
            zp = self.zero_points(band)
        m = -2.5 * np.log10(flux) + zp
        return m

    def mag_to_flux(self, mag, band, zp=None):
        if zp is None:
            zp = self.zero_points(band)
        return np.power(10., -0.4 * (mag-zp))
        
    def flux_variance(self, exptime, flux, skyflux, seeing, band):
        """The variance of PSF fluxes in the requested bands.

        Args:
          exptime (float or ndarray of floats): exposure times
          flux (float or ndarray of float, optional): source flux on the focal plane (in e/s)
          skyflux (float or ndarray of float, optional): sly brightness (in e/s/arcsecond^2)
          seeing (float or ndarray of floats): PSF FWHM in arcseconds
          band (str or ndarray of str): band

        Returns:
           float or ndarray of floats: flux variance

        .. note :: magnitude system specified in class constructor
        """
        exptime, flux, skyflux, seeing, band = self._broadcast(exptime, flux, skyflux, seeing, band)
        flx = flux * exptime
        flx[flx<0.] = 1.E-6
        v = self.etc.flux_variance(flx, # flux*exptime,
                                   skyflux*exptime*self.pixel_area,
                                   seeing/self.pixel_size,
                                   ron=self.ron)
        return v / exptime**2

    def signal_to_noise(self, exptime, flux, skyflux, seeing, band):
        """The SNR ratios 
        """
        exptime, flux, skyflux, seeing, band = self._broadcast(exptime, flux, skyflux, seeing, band)
        #        zp = np.asarray([self.zp[b] for b in band])
        #        flux = np.power(10., -0.4 * (mag-zp)) * exptime
        #        sky = np.power(10., -0.4 * (skymag-zp)) * self.pixel_area * exptime
        var = self.etc.flux_variance(flux*exptime,
                                     skyflux*exptime*self.pixel_area,
                                     seeing/self.pixel_size,
                                     ron=self.ron)
        return flux / np.sqrt(var)
        
    def flux_lim(self, exptime, skyflux, seeing, band, snr=5.):
        """Limiting fluxes
        """
        exptime, skyflux, seeing, band = self._broadcast(exptime, skyflux, seeing, band)
        #        zp = np.asarray([self.zp[b] for b in band])
        #        sky = np.power(10., -0.4 * (skymag-zp)) * self.pixel_area * exptime
        flim = self.etc.flux_lim(skyflux*exptime*self.pixel_area,
                                 seeing/self.pixel_size,
                                 ron=self.ron,
                                 snr=snr)
        return flim/exptime

    def mag_lim(self, exptime, skyflux, seeing, band, snr=5.):
        """Magnitude limit
        
        """
        exptime, skyflux, seeing, band = self._broadcast(exptime, skyflux, seeing, band)
        zp = np.asarray([self.zp[b] for b in band])
        #        sky = np.power(10., -0.4 * (skymag-zp)) * self.pixel_area * exptime
        flux_lim = self.etc.flux_lim(skyflux*exptime*self.pixel_area, 
                                     seeing/self.pixel_size, 
                                     ron=self.ron, 
                                     snr=snr)
        return -2.5 * np.log10(flux_lim/exptime) + zp
    
    def dump(self, exptime=3600., return_data=False):
        """
        """
        d = self.data
        nea = [self.etc.noise_equivalent_area(s/self.pixel_size) for s in self.data['iq']]
        exptime = [exptime] * len(d)
        skyflux = self.mag_to_flux(d['mag_sky'], d['band'])
        m5 = [float(self.mag_lim(t,m,s,b,snr=5.)) \
                  for t,m,s,b in zip(exptime, skyflux, d['iq'], d['band'])]
        
        print(('-' * 72))
        s = '%-10s  %8s %6s   %6s  %6s  %6s  %6s  %6s %6s' % ('band', 'leff',  'ZP',   'msky',    'fsky',  'FWHM',   'NEA',    'gFWHM', 'm5')
        print(s)
        s = '%-10s  %8s %6s   %6s  %6s  %-6s  %6s  %6s %6s' % ('',     '(AA)',  '(AB)', '(m/"^2)', '(e/px/s)', '(")', '(px^2)', '(")', '')
        print(s)
        print(('-' * 72))
        for i,r in enumerate(self.data):
            s = '%-10s  %8.2f  %6.2f   %6.2f  %6.2f %6.2f  %6.2f  %6.2f %6.2f'  % \
                (r['band'], r['leff'], r['zp'], r['mag_sky'], r['flux_sky'], r['iq'], nea[i], 0., m5[i])
            print(s)
        print(('-' * 72))
        
        if return_data:
            return self.data, np.array(nea), np.array(m5)
        
    def plot(self, wl_range=None):
        import pylab as pl 
        r = self.data['leff'].min()-2000., self.data['leff'].max()+2000.
        wl = np.arange(r[0], r[1], 10.)
        color = [pl.cm.jet(int((x-r[0])/(r[1]-r[0]) * 256.)) for x in self.data['leff']]        
        ma = self.instrument.area
        t = self.instrument.mirror_reflectivity
        pl.plot(wl, t(wl), ls='-', lw=2, color='gray', label='$M(\lambda)$')
        t = t * self.instrument.optics_transmission
        pl.plot(wl, t(wl), ls='--', color='gray', label='$M(\lambda) O(\lambda)$')
        t = t * self.instrument.atm_trans
        pl.plot(wl, t(wl), ls='-.', color='gray', label='$M(\lambda) O(\lambda) A(\lambda)$')
        t = t * self.instrument.ccd_qe
        pl.plot(wl, t(wl), ls='-', color='gray', label='$M(\lambda) O(\lambda) A(\lambda) \epsilon(\lambda)$')
        for i,r in enumerate(self.data):
            pl.plot(wl, self.data['tr'][i](wl)/ma, ls='-', color=color[i])
        if wl_range is not None:
            pl.xlim(wl_range)
        pl.ylim((0., 1.2))
        pl.xlabel('$\lambda\ \mathrm{[nm]}$')
        pl.ylabel('$T(\lambda)$')
        pl.legend(loc=9, ncol=2)
    
    def save(self, filename):
        pass

    def load(self, filename):
        pass


_instrument_repository = {}

def find(instrument_name):
    """
    Search for the requested instrument in a global dictionary.  If
    not found, load the instrument, return it (and update the cache).
    """
    global _instrument_repository
    ret = _instrument_repository.get(instrument_name, None)
    if ret is None:
        ret = ImagingInstrument(instrument_name)
        _instrument_repository[instrument_name] = ret
    return ret

def band_colors(band_names):
    import pylab as pl
    blist = np.unique(band_names)
    instrument_names = np.asarray([r[0] for r in np.core.defchararray.split(blist, b':')])
    instruments = [find(nm) for nm in instrument_names]
    data = np.hstack([instr.data for instr in instruments])
    color = [pl.cm.jet(int((wl-3000) * 256. /(11000-3000.))) for wl in data['leff']]
    return dict(list(zip(data['band'], color)))

def filters(fn):
    pass
    

def polyfit2d(x, y, f, deg):
    """Fit a 2D polynomial on experimental values. 
    
    Given points (x,y) and corresponding measurements (f) fit the 2D
    polynomial of degree deg=(degx, degy). 
    
    Args:
        x,y (ndarrays of floats): positions on the plane 
        f (ndarray of float): measurements
        deg (tuple): degree of the polynomial to fit
            deg = (x-degree, y-degree)
       
    Returns:
        the least-square solution, reshaped so that it is 
        directly usable by polyval2d        
    """
    x = np.asarray(x)
    y = np.asarray(y)
    f = np.asarray(f)
    deg = np.asarray(deg)
    J = polyvander2d(x, y, deg)
    J = J.reshape(-1, J.shape[-1])
    f = f.reshape((J.shape[0],))
    c = np.linalg.lstsq(J,f)[0]
    return c.reshape(deg+1)


        # zp = None
        # if mag is not None and flux is None:
        #     zp = np.asarray([self.zp[b] for b in band])
        #     flux = np.power(10., -0.4 * (mag-zp)) * exptime
        # elif mag is None and flux is not None:
        #     flux *= exptime
        # else:
        #     raise ValueError('specify mag or flux')
        # if skymag is not None and skyflux is None:
        #     if not zp:
        #         zp = np.asarray([self.zp[b] for b in band])
        #     sky = np.power(10., -0.4 * (skymag-zp)) * self.pixel_area * exptime
        # elif skymag is None and skyflux is not None:
        #     sky *= self.pixel_area * exptime
        # else:
        #     raise ValueError('specify either sky mag or sky flux')
        # return self.etc.flux_variance(flux, sky, seeing/self.pixel_size, ron=self.ron)
