class Registry(object):
    """Connect strings to instances and loaders."""

    def __init__(self):
        self._loaders = OrderedDict()
        self._instances = OrderedDict()
        self._primary_loaders = []  # keys of _loaders not including aliases

    def register_loader(self, name, func, args=None, version=None, meta=None,
                        force=False):
        """Register a data reading function.

        Parameters
        ----------
        name : str
            The data identifier.
        func : callable
            The function to read in the data. Must accept a name and version
            keyword argument.
        args : list, optional
            Arguments to pass to the function. Default is an empty list.
        version : str, optional
            Sub-version of name, if desired. Use formats such as ``'1'``,
            ``'1.0'``, ``'1.0.0'``, etc. Default is `None`.
        force : bool, optional
            Whether to override any existing function if already present.
        meta : dict, optional
            Metadata describing this loader. Default is an empty dictionary.
        """

        if args is None:
            args = []
        if meta is None:
            meta = {}

        name = name.lower()  # names are stored in all-lowercase.
        key = (name, version)

        # check if key already exists in the registry
        if key in self._loaders and not force:
            versionstr = \
                "" if version is None else " (version={!r})".format(version)
            raise Exception("Loader named {!r}{:s} is already "
                            "defined. Use force=True to override."
                            .format(name, versionstr))

        self._loaders[key] = func, args, meta
        self._primary_loaders.append(key)

    def register(self, instance, name=None, force=False):
        """Register a class instance.

        Parameters
        ----------
        instance : object
            The object to be registered.
        name : str, optional
            Identifier. If `None`, the name is taken from the `name` attribute
            of the instance, if it exists and is a string.
        force : bool, optional
            Whether to override any existing instance of the same name.
            Note: this may not play well with versioned instances.
        """

        if name is None:
            try:
                name = instance.name
            except AttributeError:
                raise ValueError("name not given and instance has no 'name' "
                                 "attribute")
            if not isinstance(name, str):
                raise ValueError("name attribute of {0!r:s} is not a string.")

        name = name.lower()
        key = (name, None)

        # check if key is already in instances or loaders:
        if (key in self._instances or key in self._loaders) and not force:
            raise Exception("{0:s} already in registry. Use force=True"
                            " to override.".format(name))

        self._instances[key] = instance

    def alias(self, new_name, existing_name, new_version=None,
              existing_version=None):
        """Alias a new name to an existing name."""

        found = False

        new_key = (new_name, new_version)
        existing_key = (existing_name, existing_version)

        if existing_key in self._loaders:
            found = True
            self._loaders[new_key] = self._loaders[existing_key]

        if existing_key in self._instances:
            found = True
            self._instances[new_key] = self._instances[existing_key]

        if not found:
            raise Exception("{0!r} not found in registry"
                            .format(existing_name))

    def retrieve(self, name, version=None):
        """Retrieve an instance from a registered identifier.

        Parameters
        ----------
        name : str
            Identifier of the specific instance. `name` is case-independent
            (internally names are stored as lowercase).
        version : str
            Sub-identifier. If `None`, default to highest or only version.

        Returns
        -------
        instance

        Notes
        -----
        **Precedence** The following are tried in this order:

        1. If ``name`` is already loaded in the registry, that instance is
           returned.
        2. If there is a loader defined for ``name``, it is used
           to create an instance, save it to the registry and return it.
        3. An Exception is raised listing the available registered names.

        **Versioning** There is support for multiple versions of data for the
        same `name`.

        1. If ``version`` is specified, the registry and its loaders are
           searched for ``(name, version)``.
        2. If ``version`` is not specified but there are registered loaders for
           ``(name, version)``, the latest version is used, and both
           ``(name, None)`` and ``(name, version)`` are saved to the registry.
           "Latest" is defined by string comparision.

        """

        name = name.lower()
        key = (name, version)

        # Try to retrieve from instances
        try:
            return self._instances[key]
        except KeyError:
            pass

        # Try to retrieve from the loaders.
        if key in self._loaders:
            func, args, meta = self._loaders[key]
            if version is None:
                self._instances[key] = func(*args, name=name)
            else:
                self._instances[key] = func(*args, name=name, version=version)
            return self._instances[key]

        # If we got this far and the version is not specified,
        # find the latest version and try to load it.
        if version is None:
            latest_version = None
            for regkey in self._loaders.keys():
                if (name == regkey[0] and (latest_version is None or
                                           regkey[1] > latest_version)):
                    latest_version = regkey[1]
            if latest_version is not None:
                regkey = (name, latest_version)
                func, args, meta = self._loaders[regkey]
                self._instances[regkey] = func(*args, name=name,
                                               version=latest_version)
                self._instances[key] = self._instances[regkey]
                return self._instances[key]

        # At this point we will raise an exception and all the following
        # is to make it more informative.

        # all names regardless of version:
        registered_names = set(k[0] for k in self._loaders)
        registered_names.update(set(k[0] for k in self._instances))

        # If we don't have the name regardless of version:
        if version is None or name not in registered_names:
            raise Exception(
                "{0!r} not in registry. Registered names: '{1}'"
                .format(name, "', '".join(registered_names)))

        # If version was specified but we don't have that specific version:
        registered_versions = [regkey[1] for regkey in self._loaders
                               if name == regkey[0]]
        raise Exception(
            "No {0!r} with version='{1:s}' in registry. Registered"
            " versions: '{2:s}'".format(name, version,
                                        "', '".join(registered_versions)))

    def get_loaders_metadata(self):
        """Return the metadata of all registered loaders.

        Returns
        -------
        loadermeta : list of dict
            Each item in the list is a dictionary containing a 'name'
            keyword, a 'version' keyword (if applicable), and the metadata
            keywords for the given loader.
        """

        result = []
        for key in self._primary_loaders:
            loader = self._loaders[key]
            name, version = key
            m = {'name': name}
            if version is not None:
                m['version'] = version
            m.update(loader[2])
            result.append(m)

        return result class Bandpass(object):
    """Transmission as a function of spectral wavelength.

    Parameters
    ----------
    wave : list_like
        Wavelength. Monotonically increasing values.
    trans : list_like
        Transmission fraction.
    wave_unit : `~astropy.units.Unit` or str, optional
        Wavelength unit. Default is Angstroms.
    trans_unit : `~astropy.units.Unit`, optional
        Transmission unit. Can be `~astropy.units.dimensionless_unscaled`,
        indicating a ratio of transmitted to incident photons, or units
        proportional to inverse energy, indicating a ratio of transmitted
        photons to incident energy. Default is ratio of transmitted to
        incident photons.
    normalize : bool, optional
        If True, normalize fractional transmission to be 1.0 at peak.
        It is recommended to set normalize=True if transmission is in units
        of inverse energy. (When transmission is given in these units, the
        absolute value is usually not significant; normalizing gives more
        reasonable transmission values.) Default is False.
    trim_level : float, optional
        If given, crop bandpass to region where transmission is above this
        fraction of the maximum transmission. For example, if maximum
        transmission is 0.5, ``trim_level=0.001`` will remove regions where
        transmission is below 0.0005. Only contiguous regions on the sides
        of the bandpass are removed.
    name : str, optional
        Identifier. Default is `None`.

    Examples
    --------
    Construct a Bandpass and access the input arrays:

    >>> b = Bandpass([4000., 4200., 4400.], [0.5, 1.0, 0.5])
    >>> b.wave
    array([ 4000.,  4200.,  4400.])
    >>> b.trans
    array([ 0.5,  1. ,  0.5])

    Bandpasses act like continuous 1-d functions (linear interpolation is
    used):

    >>> b([4100., 4300.])
    array([ 0.75,  0.75])

    The effective (transmission-weighted) wavelength is a property:

    >>> b.wave_eff
    4200.0

    The ``trim_level`` keyword can be used to remove "out-of-band"
    transmission upon construction. The following example removes regions of
    the bandpass with tranmission less than 1 percent of peak:

    >>> band = Bandpass([4000., 4100., 4200., 4300., 4400., 4500.],
    ...                 [0.001, 0.002,   0.5,   0.6, 0.003, 0.001],
    ...                 trim_level=0.01)

    >>> band.wave
    array([ 4100.,  4200.,  4300.,  4400.])

    >>> band.trans
    array([ 0.002,  0.5  ,  0.6  ,  0.003])

    While less strictly correct than including the "out-of-band" transmission,
    only considering the region of the bandpass where transmission is
    significant can improve model-bandpass overlap as well as performance.
    """

    def __init__(self, wave, trans, wave_unit=u.AA,
                 trans_unit=u.dimensionless_unscaled, normalize=False,
                 name=None, trim_level=None):
        wave = np.asarray(wave, dtype=np.float64)
        trans = np.asarray(trans, dtype=np.float64)
        if wave.shape != trans.shape:
            raise ValueError('shape of wave and trans must match')
        if wave.ndim != 1:
            raise ValueError('only 1-d arrays supported')

        # Ensure that units are actually units and not quantities, so that
        # `to` method returns a float and not a Quantity.
        wave_unit = u.Unit(wave_unit)
        trans_unit = u.Unit(trans_unit)

        if wave_unit != u.AA:
            wave = wave_unit.to(u.AA, wave, u.spectral())

        # If transmission is in units of inverse energy, convert to
        # unitless transmission:
        #
        # (transmitted photons / incident photons) =
        #      (photon energy) * (transmitted photons / incident energy)
        #
        # where photon energy = h * c / lambda
        if trans_unit != u.dimensionless_unscaled:
            trans = (HC_ERG_AA / wave) * trans_unit.to(u.erg**-1, trans)

        # Check that values are monotonically increasing.
        # We could sort them, but if this happens, it is more likely a user
        # error or faulty bandpass definition. So we leave it to the user to
        # sort them.
        if not np.all(np.ediff1d(wave) > 0.):
            raise ValueError('bandpass wavelength values must be monotonically'
                             ' increasing when supplied in wavelength or '
                             'decreasing when supplied in energy/frequency.')

        if normalize:
            trans /= np.max(trans)

        # Trim "out-of-band" transmission
        if trim_level is not None:
            s = slice_exclude_below(trans, np.max(trans) * trim_level, grow=1)
            wave = wave[s]
            trans = trans[s]

        # if more than one leading or trailing transmissions are zero, we
        # can remove them.
        if ((trans[0] == 0.0 and trans[1] == 0.0) or (trans[-1] == 0.0 and
                                                      trans[-2] == 0.0)):
            i = 0
            while i < len(trans) and trans[i] == 0.0:
                i += 1
            if i == len(trans):
                raise ValueError('all zero transmission')
            j = len(trans) - 1
            while j >= 0 and trans[j] == 0.0:
                j -= 1

            # back out to include a single zero
            if i > 0:
                i -= 1
            if j < len(trans) - 1:
                j += 1

            wave = wave[i:j+1]
            trans = trans[i:j+1]

        self.wave = wave
        self.trans = trans

        # Set up interpolation.
        # This appears to be the fastest-evaluating interpolant in
        # scipy.interpolate.
        self._tck = splrep(self.wave, self.trans, k=1)

        self.name = name

    def minwave(self):
        return self.wave[0]

    def maxwave(self):
        return self.wave[-1]

    @lazyproperty
    def wave_eff(self):
        """Effective wavelength of bandpass in Angstroms."""
        wave, _ = integration_grid(self.minwave(), self.maxwave(),
                                   SPECTRUM_BANDFLUX_SPACING)
        weights = self(wave)
        return np.sum(wave * weights) / np.sum(weights)

    def __call__(self, wave):
        return splev(wave, self._tck, ext=1)

    def __repr__(self):
        name = ''
        if self.name is not None:
            name = ' {!r}'.format(self.name)
        return "<{:s}{:s} at 0x{:x}>".format(self.__class__.__name__, name,
                                             id(self))

    def shifted(self, factor, name=None):
        """Return a new Bandpass instance with all wavelengths
        multiplied by a factor."""
        return Bandpass(factor * self.wave, self.trans, name=name) def get_bandpass(name, *args):
    """Get a Bandpass from the registry by name."""
    if isinstance(name, Bandpass):
        return name
    if len(args) == 0:
        return _BANDPASSES.retrieve(name)
    else:
        interp = _BANDPASS_INTERPOLATORS.retrieve(name)
        return interp.at(*args) HC_ERG_AA = H_ERG_S * C_AA_PER_S MODEL_BANDFLUX_SPACING = 5.0 def read_griddata_ascii(name_or_obj):
    """Read 2-d grid data from a text file.

    Each line has values `x0 x1 y`. Space separated.
    x1 values are only read for first x0 value. Others are assumed
    to match.

    Parameters
    ----------
    name_or_obj : str or file-like object

    Returns
    -------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        2-d array of shape (len(x0), len(x1)).
    """

    if isinstance(name_or_obj, str):
        f = open(name_or_obj, 'r')
    else:
        f = name_or_obj

    x0 = []    # x0 values.
    x1 = None  # x1 values for first x0 value, assume others are the same.
    y = []     # 2-d array of internal values

    x0_current = None
    x1_current = []
    y1_current = []
    for line in f:
        stripped_line = _stripcomment(line)
        if len(stripped_line) == 0:
            continue
        x0_tmp, x1_tmp, y_tmp = map(float, stripped_line.split())
        if x0_current is None:
            x0_current = x0_tmp  # Initialize first time

        # If there is a new x0 value, ingest the old one and reset values
        if x0_tmp != x0_current:
            x0.append(x0_current)
            if x1 is None:
                x1 = x1_current
            y.append(y1_current)

            x0_current = x0_tmp
            x1_current = []
            y1_current = []

        x1_current.append(x1_tmp)
        y1_current.append(y_tmp)

    # Ingest the last x0 value and y1 array
    x0.append(x0_current)
    y.append(y1_current)

    f.close()
    return np.array(x0), np.array(x1), np.array(y) def read_griddata_fits(name_or_obj, ext=0):
    """Read a multi-dimensional grid of data from a FITS file, where the
    grid coordinates are encoded in the FITS-WCS header keywords.

    Parameters
    ----------
    name_or_obj : str or file-like object

    Returns
    -------
    x0, x1, ... : `~numpy.ndarray`
        1-d arrays giving coordinates of grid. The number of these arrays will
        depend on the dimension of the data array. For example, if the data
        have two dimensions, a total of three arrays will be returned:
        ``x0, x1, y``, with ``x0`` giving the coordinates of the first axis
        of ``y``. If the data have three dimensions, a total of four arrays
        will be returned: ``x0, x1, x2, y``, and so on with higher dimensions.
    y : `~numpy.ndarray`
        n-d array of shape ``(len(x0), len(x1), ...)``. For three dimensions
        for example, the value at ``y[i, j, k]`` corresponds to coordinates
        ``(x0[i], x1[j], x2[k])``.
    """

    hdulist = fits.open(name_or_obj)
    w = wcs.WCS(hdulist[ext].header)
    y = hdulist[ext].data

    # get abcissa values (coordinates at grid values)
    xs = []
    for i in range(y.ndim):
        j = y.ndim - i  # The i-th axis (in Python) corresponds to FITS AXISj
        coords = np.zeros((y.shape[i], y.ndim), dtype=np.float32)
        coords[:, j-1] = np.arange(y.shape[i])
        x = w.wcs_pix2world(coords, 0)[:, j-1]
        xs.append(x)

    hdulist.close()

    return tuple(xs) + (y,) def read_multivector_griddata_ascii(name_or_obj):
    """Read 2-d grid data from a text file.

    Each line has values `x0 x1 y0 y1 ...`. Space separated.
    Assumed to be grid of values.

    Parameters
    ----------
    name_or_obj : str or file-like object
        The name of the file or a file-like object containing the
        data.

    Returns
    -------
    x0 : numpy.ndarray
        1-d array.
    x1 : numpy.ndarray
        1-d array.
    y : numpy.ndarray
        3-d array of shape ``(n, len(x0), len(x1))`` where ``n`` is
        the number of y values on each line.
    """
    data = np.loadtxt(name_or_obj)

    x0 = np.sort(np.unique(data[:, 0]))
    x1 = np.sort(np.unique(data[:, 1]))
    y = np.zeros((len(data[0]) - 2, len(x0), len(x1)))

    for i0, p in enumerate(x0):
        for i1, q in enumerate(x1):
            ind = (data[:, 0] == p) & (data[:, 1] == q)
            y[:, i0, i1] = data[ind, 2:]

    return x0, x1, y def get_magsystem(name):
    """Get a MagSystem from the registry by name."""
    if isinstance(name, MagSystem):
        return name
    return _MAGSYSTEMS.retrieve(name) def integration_grid(low, high, target_spacing):
    """Divide the range between `start` and `stop` into uniform bins
    with spacing less than or equal to `target_spacing` and return the
    bin midpoints and the actual spacing."""

    range_diff = high - low
    spacing = range_diff / int(math.ceil(range_diff / target_spacing))
    grid = np.arange(low + 0.5 * spacing, high, spacing)

    return grid, spacing
import abc
import os
import extinction
import numpy as np
            import scipy  # to get scipy version
from copy import copy as cp
from math import ceil
from textwrap import dedent
from astropy import (cosmology, units as u)
from astropy.utils.misc import isiterable
from scipy.interpolate import (
    InterpolatedUnivariateSpline as Spline1d,
    RectBivariateSpline as Spline2d
)
from ._registry import Registry
from .bandpasses import Bandpass, get_bandpass
from .constants import HC_ERG_AA, MODEL_BANDFLUX_SPACING
from .io import (
    read_griddata_ascii, read_griddata_fits,
    read_multivector_griddata_ascii
)
from .magsystems import get_magsystem
from .salt2utils import BicubicInterpolator, SALT2ColorLaw
from .utils import integration_grid
            from scipy.interpolate import RegularGridInterpolator
# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""Classes to represent spectral models of astronomical transients."""




__all__ = ['get_source', 'Source', 'TimeSeriesSource', 'StretchSource',
           'SUGARSource', 'SALT2Source', 'SALT3Source', 'MLCS2k2Source',
           'SNEMOSource', 'Model', 'PropagationEffect', 'CCM89Dust',
           'OD94Dust', 'F99Dust']

_SOURCES = Registry()


def _check_for_fitpack_error(e, a, name):
    """Raise a more informative error message for fitpack errors.

    This is implemented as a separate function rather than subclassing Spline2d
    so that we can raise the error closer to user-facing functions. We
    may wish to change this behavior in the future. For example, if some
    models are implemented not based on RectBivariateSpline so that they
    don't have this restriction.

    Parameters
    ----------
    e : ValueError
    a : `~numpy.ndarray` (0-d or 1-d)
    name : str
    """

    # Check if the error is a specific one raise by RectBivariateSpline
    # If it is, check if supplied array is *not* monotonically increasing
    if (len(e.args) > 0 and
            e.args[0].startswith("Error code returned by bispev: 10") and
            np.any(np.ediff1d(a) < 0.)):
        raise ValueError(name + ' must be monotonically increasing')


def get_source(name, version=None, copy=False):
    """Retrieve a Source from the registry by name.

    Parameters
    ----------
    name : str
        Name of source in the registry.
    version : str, optional
        Version identifier for sources with multiple versions. Default is
        `None` which corresponds to the latest, or only, version.
    copy : bool, optional
        If True and if `name` is already a Source instance, return a copy of
        it. (If `name` is a str a copy of the instance
        in the registry is always returned, regardless of the value of this
        parameter.) Default is False.
    """

    # If we need to retrieve from the registry, we want to return a shallow
    # copy, in order to keep the copy in the registry "pristene". However, we
    # *don't* want a shallow copy otherwise. Therefore,
    # we need to check if `name` is already an instance of Model before
    # going to the registry, so we know whether or not to make a shallow copy.
    if isinstance(name, Source):
        if copy:
            return cp(name)
        else:
            return name
    else:
        return cp(_SOURCES.retrieve(name, version=version))


def _bandflux_single(model, band, time_or_phase):
    """Synthetic photometry of model through a single bandpass.

    Parameters
    ----------
    model : Source or Model
    band : Bandpass
    time_or_phase : `~numpy.ndarray` (1-d)
    """

    # Check that bandpass wavelength range is fully contained in model
    # wavelength range.
    if (band.minwave() < model.minwave() or band.maxwave() > model.maxwave()):
        raise ValueError('bandpass {0!r:s} [{1:.6g}, .., {2:.6g}] '
                         'outside spectral range [{3:.6g}, .., {4:.6g}]'
                         .format(band.name, band.minwave(), band.maxwave(),
                                 model.minwave(), model.maxwave()))

    # Set up wavelength grid. Spacing (dwave) evenly divides the bandpass,
    # closest to 5 angstroms without going over.
    wave, dwave = integration_grid(band.minwave(), band.maxwave(),
                                   MODEL_BANDFLUX_SPACING)
    trans = band(wave)
    f = model._flux(time_or_phase, wave)


    return np.sum(wave * trans * f, axis=1) * dwave / HC_ERG_AA