#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Identifies events of small particles crossing through an ion channel.

This script parses a trajectory and of an ion-channel-in-membrane simulation
and identifies events of small target particles (e.g. ions, water molecules)
crossing the ion channel. A crossing event is defined as an ion moving from the
bulk domain on one side of the membrane through the channel pore domain to the
bulk domain on the opposite side of the membrane. Cases where the target
particle moves through the periodic boundary or through the membrane itself are
not classified as crossing events.

Particles are assigned to the two bulk domains based on whether their
z-coordinate is smaller or larger than the mean z-coordinate of the membrane.
However, if a particle is located within a cylindrical region around the centre
of geometry of the channel protein, it is assigned to the pore domain instead.
Crossing events are then identified based on the particles moving between these
three domains.

The user can specify an MDAnalysis selection that defines which particles will
be considered. It should be noted that the script operates on a per-atom basis.
Therefore, if the user selected e.g. water molecules on a per-residue bases, the
crossing of one water molecule will be recorded as the three individual
crossings of the constituent atoms.

Output is written to a JSON file and corresponds to a data frame structure in
which each row represents one crossing event. The columns will contain the
(res)ID, (res)name, charge, and mass of the particle that crossed the pore
alongside the times at which the particle enters and exits the pore. A crossing
number is given which indicates the direction in which the particle crossed
the pore (i.e. +1 for crossing in positive z-direction and -1 for crossing in
negative z-direction).

The overall number of particles that cross the pore can be readily calculated
from this data as the cumulative sum over the crossing number. Similarly, the
total amount of charge or mass transported can be evaluated as the cumulative
sum over the product of the crossing number with charge and mass respectively.
"""

import argparse

import numpy as np
import pandas as pd

import MDAnalysis as mda

def determine_domain_indices(
    topology_file,
    trajectory_file,
    sel_pore,
    sel_bilayer,
    sel_target,
    r_margin,
    z_margin,
    b,
    e,
    dt
):
    """Passes through trajectory and assigns domain indices.

    The return value is a data frame containing three columns for particle ID,
    time stamp, and the domain index of that particle at the given time. The
    domain index gives an indication of whether a particle is located below
    the membrane (-1), above the membrane (+1), or inside the channel pore (0).
    The user can set a start time, end time, and time step based on which frames
    will be analysed or skipped.

    Particles are assigned to the -1 and +1 domains based on whether their
    z-coordinate is larger or smaller than the mean z-coordinate of the
    membrane. If a particle is located within a cylindrical region around the
    center of geometry of the protein, it is assigned the domain index 0
    instead. Radius and z-extend of the cylindrical region are determined from
    the bounding box of the protein, but the user may add a margin to either
    of these parameters.
    """

    # create MDA universe:
    u = mda.Universe(topology_file, trajectory_file)

    # determine atom groups for analysis:
    protein = u.select_atoms(sel_pore)
    membrane = u.select_atoms(sel_bilayer)
    target = u.select_atoms(sel_target)

    # find pore radius:
    prot_radius = max(np.diff(protein.bbox(), axis=0)[0][0:2]) / 2.0
    radius = prot_radius + r_margin

    # create data frame of target positions over time:
    dat = []
    for ts in u.trajectory:

        # skip unwanted frames:
        if ts.time < b or ts.time > e or ts.time % dt != 0.0:
            continue

        # inform user:
        print(
            "  ~> identifying domain indices in frame: " + str(ts.frame)
            + " at time: " + str(ts.time)
        )

        # create data frame from target particle positions:
        tmp = pd.DataFrame(
            target.positions,
            columns=["x", "y", "z"]
        )

        # add particle properties:
        tmp["particle_id"] = target.ids
        tmp["name"] = [str(x) for x in target.names]
        tmp["resname"] = [str(x) for x in target.resnames]
        tmp["resid"] = [int(x) for x in target.resids]
        tmp["charge"] = target.charges
        tmp["mass"] = target.masses

        # add time stamp:
        tmp["t"] = ts.time

        # position of protein COG in x/y-plane:
        x_cen = protein.center_of_geometry()[0]
        y_cen = protein.center_of_geometry()[1]

        # z-extent of proteina nd middle of membrane:
        z_min = protein.bbox()[0, 2] - z_margin
        z_max = protein.bbox()[1, 2] + z_margin
        z_mid = np.mean(membrane.bbox(), axis=0)[2]

        # assign domain indices:
        tmp.loc[
            (tmp["z"] < z_mid),
            "domain_idx"
        ] = int(-1)
        tmp.loc[
            (tmp["z"] >= z_mid),
            "domain_idx"
        ] = int(1)
        tmp.loc[
            (tmp["z"] >= z_min)
            & (tmp["z"] <= z_max)
            & ((tmp["x"] - x_cen)**2 + (tmp["y"] - y_cen)**2 < radius**2),
            "domain_idx"
        ] = int(0)

        # sanity check:
        num_nan = np.sum(tmp["domain_idx"].isna())
        if num_nan > 0:
            raise Exception("Could not unambiguously assign domain indices.")

        # drop unneccessary columns:
        tmp = tmp.drop(columns=["x", "y", "z"])

        # add to data frame list:
        dat.append(tmp)

    # combine into overall data frame:
    df = pd.concat(dat)

    # return the overall data frame of domain indices:
    return(df)

def aggregate_event_pairs(df):
    """Aggregates pairs of subsequent events.

    The resulting data frame will contain as many rows as the number of
    crossing events (i.e. a particle going through the channel) occuring in the
    given data frame. This means that the output data frame will be empty if no
    crossing event occurs. Each row will contain the time at which the particle
    entered and left the pore region as well as the mean of those two times. In
    addition it will contain a crossing number, which indicates whether the
    particle crossed the pore in positive or negative z-direction (its absolute
    value is always one).

    The input data frame must contain the domain index difference for a single
    particle only and may not contain any jump events, i.e. events where the
    particle moved from domain +1 to domain -1 (or vice versa) without passing
    through domain 0 (the pore). This can happen if the particle jumps through
    the periodic boundary or if it moves through the membrane.

    The crossing number is calculated as the mean of two subsequent domain index
    differences. If they have opposite sign, the crossing number will be zero
    and the row will be dropped. If they have the same sign, this indicates that
    the particle moved either from domain -1 to domain 0 to domain +1 or from
    domain +1 to domain 0 to domain -1, corresponding to a true crossing event.
    The sign of the crossing number will reflect the direction of the crossing.
    """

    # make copy of input data frame to avoid overwriting contents:
    df = df.copy()

    # make sure data frame has even number of rows:
    if df.shape[0] % 2 is not 0:
        df = df.iloc[:-1]

    # introduce pairing index:
    df["pairing"] = np.repeat(range(0, int(df.shape[0]/2)), 2)

    # calculate summry statistics over pairings:
    df["crossing_number"] = df.groupby("pairing").domain_diff.transform("mean")
    df["t_enter"] = df.groupby("pairing").t.transform("min")
    df["t_mean"] = df.groupby("pairing").t.transform("mean")
    df["t_exit"] = df.groupby("pairing").t.transform("max")
    df = df.groupby(["pairing", "name"]).mean()

    # remove all rows not corresponding to a crossing event:
    df = df.loc[df.domain_diff != 0.0, ]

    # return data frame to caller:
    return(df)

def identify_crossing_events(df):
    """Identifies pore crossing events in data frame of domain indices.

    Given a data frame of domain indices alongside time stamps and particle IDs,
    this function identifies crossing events, i.e. a particle moving from one
    side of the membrane to the other THROUGH THE PORE REGION. Cases where a
    particle moves through the periodic boundary or through the membrane itself
    (i.e. outside the cylindrical protein region) are not counted as crossing
    events.

    The return value is a data frame in which each row corresponds to a
    crossing event and the columns identify times at which the particle entered
    and exited the pore, the mean of these times, and the ID of the particle
    which crosses the pore. Note that the input data frame should contain only
    a single unique particle ID, i.e. this function should by run on the domain
    index data frame grouped by particle indices. In addition, a further column
    contains a crossing number indicating the direction in which the particle
    passed the pore.

    This is done by first calculating the difference between subsequent domain
    indices. Where this difference is zero, the particle stayed in the given
    domain and the corresponding row can be ignored. Where this difference is
    +1/-1, the particle moved into or out of the pore. Where the difference is
    +2/-2, the particle moved across the periodic boundary or through the
    membrane.

    The time series of domain index differences is then partitioned into
    individual series between periodic boundary jump events. In each such
    subdivision the target particle is known to start outside the pore domain.
    A crossing event can then be identified by averaging over the sum of
    pairs of subsequent domain index differences. Note that for a particle that
    starts out inside the pore domain at the beginning of the trajectory no
    clear crossing direction can be established (i.e. is it going back to the
    domain it came from or is it moving across the pore to the opposite domain).
    By convention, these events will therefore not be counted as pore crossing
    events.
    """

    # make copy of input data frame to avoid overwriting contents:
    df = df.copy()

    # calculate difference between subsequent indices:
    df["domain_diff"] = df["domain_idx"].shift(0) - df["domain_idx"].shift(1)
    df["t"] = 0.5*(df["t"].shift(0) + df["t"].shift(1))
    df = df.dropna()
    df = df.drop(columns=["domain_idx"])

    # drop cases where particle did not change domain at all:
    df = df.loc[np.abs(df["domain_diff"]) != 0.0, ]

    # create index for number of jumps across period boundary (or through the
    # membrane)
    df["pbcjump"] = 0
    df.pbcjump[(df.domain_diff == -2) | (df.domain_diff == 2)] = 1
    df.pbcjump = df.pbcjump.cumsum()

    # drop all pbcjump groupings with less then two non-pbcjump events:
    df = df.loc[(df.domain_diff != -2) & (df.domain_diff != 2), ]
    df["pbccount"] = df.groupby("pbcjump").domain_diff.transform("count")
    df = df.loc[df.pbccount > 1, ]
    df = df.drop(columns=["pbccount"])

    # handle special case of empty data frame:
    if df.empty is True:

        # manually add columns:
        df["crossing_number"] = None
        df["t_enter"] = None
        df["t_mean"] = None
        df["t_exit"] = None

    else:

        # aggregate pairwise domain index differences:
        df = df.groupby("pbcjump").apply(aggregate_event_pairs)

    # remove unneccessary columns and indices:
    df = df.drop(columns=["t", "domain_diff", "pbcjump"])
    df = df.reset_index(drop=True)

    # return data frame:
    return(df)

def main(argdict):
    """Main function for entry point checking.

    Expects to be given a dictionary of command line arguments.
    """

    # parse trajectory and assign domain indices:
    print("--> identifying domain indices...")
    nm2ang = 10.0
    domain_indices = determine_domain_indices(
        argdict["s"],
        argdict["f"],
        argdict["sel_pore"],
        argdict["sel_bilayer"],
        argdict["sel_target"],
        argdict["r_margin"] * nm2ang,
        argdict["z_margin"] * nm2ang,
        argdict["b"],
        argdict["e"],
        argdict["dt"]
    )

    # find crossing events by particle:
    print("--> identifying crossing events...")
    crossing_events = (
        domain_indices.groupby(["particle_id"]).apply(identify_crossing_events)
    )
    crossing_events = crossing_events.reset_index(drop=True)

    # restore atom and residue names:
    name = domain_indices.groupby("particle_id").name.first()
    crossing_events["name"] = np.array(name.loc[crossing_events.particle_id])
    crossing_events["resname"] = np.array(name.loc[crossing_events.particle_id])

    # serialize results:
    crossing_events.to_json(argdict["o"], orient="records")
    with open(argdict["o"], "a") as f:
        f.write("\n")

# entry point check:
if __name__ == "__main__":

    # turn off this warning:
    pd.options.mode.chained_assignment = None

    # parse command line arguments:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "-s",
        type=str,
        nargs=None,
        default="production.tpr",
        help="""Topology file. Should contain charges and masses!"""
    )
    parser.add_argument(
        "-f",
        type=str,
        nargs=None,
        default="production.xtc",
        help="""Trajectory file name."""
    )
    parser.add_argument(
        "-o",
        type=str,
        nargs=None,
        default="crossing_events.json",
        help="""Name of output JSON file containing all crossing events."""
    )
    parser.add_argument(
        "-sel_target",
        type=str,
        nargs=None,
        default="resname NA CL",
        help="""MDAnalysis selection identifying the mobile particles
        crossing the pore."""
    )
    parser.add_argument(
        "-sel_pore",
        type=str,
        nargs=None,
        default="protein",
        help="""MDAnalysis selection identifying the channel domain."""
    )
    parser.add_argument(
        "-sel_bilayer",
        type=str,
        nargs=None,
        default="resname DOPC POPC",
        help="""MDAnalysis selection identifying the lipid bilayer."""
    )
    parser.add_argument(
        "-r_margin",
        type=float,
        nargs=None,
        default=0.0,
        help="""Radius margin around protein for identification of pore
        domain in nanometers. The pore domain is taken to be a cylinder whose
        radius is the radius of the protein (in the x/y-plane) plus this
        margin."""
    )
    parser.add_argument(
        "-z_margin",
        type=float,
        nargs=None,
        default=-1.0,
        help="""Margin for extension of pore domain around protein in
        nanometers. The pore domain is taken to be a cylinder whose the
        extension is determined from the z-extension of the protein plus this
        margin. Should never exceed half the length of the pore."""
    )
    parser.add_argument(
        "-b",
        type=float,
        nargs=None,
        default=0.0,
        help="""Start time for trajectory analysis in picoseconds. Frames
        before this time will be ignored."""
    )
    parser.add_argument(
        "-e",
        type=float,
        nargs=None,
        default=float("inf"),
        help="""End time for trajectory analysis in picoseconds. Frames
        after this time will be ignored."""
    )
    parser.add_argument(
        "-dt",
        type=float,
        nargs=None,
        default=100.0,
        help="""Time step for trajectroy analysis. Frames not that are not
        an integer multiple of this time step will be ignored. A time step of
        100 picoseconds seems to be sufficiently small for ions."""
    )

    # parse arguments and convert to dictionary:
    args = parser.parse_args()
    argdict = vars(args)

    # pass arguments to main function:
    main(argdict)


import warnings

import numpy as np
import operator

from .. import types, utils, config
from .templates import (AttributeTemplate, AbstractTemplate, CallableTemplate,
                        Registry, signature)

from ..numpy_support import (ufunc_find_matching_loop,
                             supported_ufunc_loop, as_dtype,
                             from_dtype, as_dtype, resolve_output_type,
                             carray, farray)
from ..errors import TypingError, NumbaPerformanceWarning
from numba import pndindex

registry = Registry()
infer = registry.register
infer_global = registry.register_global
infer_getattr = registry.register_attr

class Numpy_rules_ufunc(AbstractTemplate):
    @classmethod
    def _handle_inputs(cls, ufunc, args, kws):
        """
        Process argument types to a given *ufunc*.
        Returns a (base types, explicit outputs, ndims, layout) tuple where:
        - `base types` is a tuple of scalar types for each input
        - `explicit outputs` is a tuple of explicit output types (arrays)
        - `ndims` is the number of dimensions of the loop and also of
          any outputs, explicit or implicit
        - `layout` is the layout for any implicit output to be allocated
        """
        nin = ufunc.nin
        nout = ufunc.nout
        nargs = ufunc.nargs

        # preconditions
        assert nargs == nin + nout

        if nout > 1:
            msg = "ufunc '{0}': not supported in this mode (more than 1 output)"
            raise TypingError(msg=msg.format(ufunc.__name__))

        if len(args) < nin:
            msg = "ufunc '{0}': not enough arguments ({1} found, {2} required)"
            raise TypingError(msg=msg.format(ufunc.__name__, len(args), nin))

        if len(args) > nargs:
            msg = "ufunc '{0}': too many arguments ({1} found, {2} maximum)"
            raise TypingError(msg=msg.format(ufunc.__name__, len(args), nargs))

        args = [a.as_array if isinstance(a, types.ArrayCompatible) else a
                for a in args]
        arg_ndims = [a.ndim if isinstance(a, types.ArrayCompatible) else 0
                     for a in args]
        ndims = max(arg_ndims)

        # explicit outputs must be arrays (no explicit scalar return values supported)
        explicit_outputs = args[nin:]

        # all the explicit outputs must match the number max number of dimensions
        if not all(d == ndims for d in arg_ndims[nin:]):
            msg = "ufunc '{0}' called with unsuitable explicit output arrays."
            raise TypingError(msg=msg.format(ufunc.__name__))

        if not all(isinstance(output, types.ArrayCompatible)
                   for output in explicit_outputs):
            msg = "ufunc '{0}' called with an explicit output that is not an array"
            raise TypingError(msg=msg.format(ufunc.__name__))

        if not all(output.mutable for output in explicit_outputs):
            msg = "ufunc '{0}' called with an explicit output that is read-only"
            raise TypingError(msg=msg.format(ufunc.__name__))

        # find the kernel to use, based only in the input types (as does NumPy)
        base_types = [x.dtype if isinstance(x, types.ArrayCompatible) else x
                      for x in args]

        # Figure out the output array layout, if needed.
        layout = None
        if ndims > 0 and (len(explicit_outputs) < ufunc.nout):
            layout = 'C'
            layouts = [x.layout if isinstance(x, types.ArrayCompatible) else ''
                       for x in args]

            # Prefer C contig if any array is C contig.
            # Next, prefer F contig.
            # Defaults to C contig if not layouts are C/F.
            if 'C' not in layouts and 'F' in layouts:
                layout = 'F'

        return base_types, explicit_outputs, ndims, layout

    @property
    def ufunc(self):
        return self.key

    def generic(self, args, kws):
        ufunc = self.ufunc
        base_types, explicit_outputs, ndims, layout = self._handle_inputs(
            ufunc, args, kws)
        ufunc_loop = ufunc_find_matching_loop(ufunc, base_types)
        if ufunc_loop is None:
            raise TypingError("can't resolve ufunc {0} for types {1}".format(ufunc.__name__, args))

        # check if all the types involved in the ufunc loop are supported in this mode
        if not supported_ufunc_loop(ufunc, ufunc_loop):
            msg = "ufunc '{0}' using the loop '{1}' not supported in this mode"
            raise TypingError(msg=msg.format(ufunc.__name__, ufunc_loop.ufunc_sig))

        # if there is any explicit output type, check that it is valid
        explicit_outputs_np = [as_dtype(tp.dtype) for tp in explicit_outputs]

        # Numpy will happily use unsafe conversions (although it will actually warn)
        if not all (np.can_cast(fromty, toty, 'unsafe') for (fromty, toty) in
                    zip(ufunc_loop.numpy_outputs, explicit_outputs_np)):
            msg = "ufunc '{0}' can't cast result to explicit result type"
            raise TypingError(msg=msg.format(ufunc.__name__))

        # A valid loop was found that is compatible. The result of type inference should
        # be based on the explicit output types, and when not available with the type given
        # by the selected NumPy loop
        out = list(explicit_outputs)
        implicit_output_count = ufunc.nout - len(explicit_outputs)
        if implicit_output_count > 0:
            # XXX this is currently wrong for datetime64 and timedelta64,
            # as ufunc_find_matching_loop() doesn't do any type inference.
            ret_tys = ufunc_loop.outputs[-implicit_output_count:]
            if ndims > 0:
                assert layout is not None
                ret_tys = [types.Array(dtype=ret_ty, ndim=ndims, layout=layout)
                           for ret_ty in ret_tys]
                ret_tys = [resolve_output_type(self.context, args, ret_ty)
                           for ret_ty in ret_tys]
            out.extend(ret_tys)

        # note: although the previous code should support multiple return values, only one
        #       is supported as of now (signature may not support more than one).
        #       there is an check enforcing only one output
        out.extend(args)
        return signature(*out)

@infer_global(operator.pos)
class UnaryPositiveArray(AbstractTemplate):
    '''Typing template class for +(array) expressions.  This operator is
    special because there is no Numpy ufunc associated with it; we
    include typing for it here (numba.typing.npydecl) because this is
    where the remaining array operators are defined.
    '''
    key = operator.pos

    def generic(self, args, kws):
        assert not kws
        if len(args) == 1 and isinstance(args[0], types.ArrayCompatible):
            arg_ty = args[0]
            return arg_ty.copy()(arg_ty)

class NumpyRulesArrayOperator(Numpy_rules_ufunc):
    _op_map = {
        operator.add: "add",
        operator.sub: "subtract",
        operator.mul: "multiply",
        operator.truediv: "true_divide",
        operator.floordiv: "floor_divide",
        operator.mod: "remainder",
        operator.pow: "power",
        operator.lshift: "left_shift",
        operator.rshift: "right_shift",
        operator.and_: "bitwise_and",
        operator.or_: "bitwise_or",
        operator.xor: "bitwise_xor",
        operator.eq: "equal",
        operator.gt: "greater",
        operator.ge: "greater_equal",
        operator.lt: "less",
        operator.le: "less_equal",
        operator.ne: "not_equal",
    }

    @property
    def ufunc(self):
        return getattr(np, self._op_map[self.key])

    @classmethod
    def install_operations(cls):
        for op, ufunc_name in cls._op_map.items():
            infer_global(op)(
                type("NumpyRulesArrayOperator_" + ufunc_name, (cls,), dict(key=op))
            )

    def generic(self, args, kws):
        '''Overloads and calls base class generic() method, returning
        None if a TypingError occurred.

        Returning None for operators is important since operators are
        heavily overloaded, and by suppressing type errors, we allow
        type inference to check other possibilities before giving up
        (particularly user-defined operators).
        '''
        try:
            sig = super(NumpyRulesArrayOperator, self).generic(args, kws)
        except TypingError:
            return None
        if sig is None:
            return None
        args = sig.args
        # Only accept at least one array argument, otherwise the operator
        # doesn't involve Numpy's ufunc machinery.
        if not any(isinstance(arg, types.ArrayCompatible)
                   for arg in args):
            return None
        return sig

_binop_map = NumpyRulesArrayOperator._op_map

class NumpyRulesInplaceArrayOperator(NumpyRulesArrayOperator):
    _op_map = {
        operator.iadd: "add",
        operator.isub: "subtract",
        operator.imul: "multiply",
        operator.itruediv: "true_divide",
        operator.ifloordiv: "floor_divide",
        operator.imod: "remainder",
        operator.ipow: "power",
        operator.ilshift: "left_shift",
        operator.irshift: "right_shift",
        operator.iand: "bitwise_and",
        operator.ior: "bitwise_or",
        operator.ixor: "bitwise_xor",
    }

    def generic(self, args, kws):
        # Type the inplace operator as if an explicit output was passed,
        # to handle type resolution correctly.
        # (for example int8[:] += int16[:] should use an int8[:] output,
        #  not int16[:])
        lhs, rhs = args
        if not isinstance(lhs, types.ArrayCompatible):
            return
        args = args + (lhs,)
        sig = super(NumpyRulesInplaceArrayOperator, self).generic(args, kws)
        # Strip off the fake explicit output
        assert len(sig.args) == 3
        real_sig = signature(sig.return_type, *sig.args[:2])
        return real_sig

class NumpyRulesUnaryArrayOperator(NumpyRulesArrayOperator):
    _op_map = {
        # Positive is a special case since there is no Numpy ufunc
        # corresponding to it (it's essentially an identity operator).
        # See UnaryPositiveArray, above.
        operator.neg: "negative",
        operator.invert: "invert",
    }

    def generic(self, args, kws):
        assert not kws
        if len(args) == 1 and isinstance(args[0], types.ArrayCompatible):
            return super(NumpyRulesUnaryArrayOperator, self).generic(args, kws)

# list of unary ufuncs to register

_math_operations = [ "add", "subtract", "multiply",
                     "logaddexp", "logaddexp2", "true_divide",
                     "floor_divide", "negative", "power",
                     "remainder", "fmod", "absolute",
                     "rint", "sign", "conjugate", "exp", "exp2",
                     "log", "log2", "log10", "expm1", "log1p",
                     "sqrt", "square", "reciprocal",
                     "divide", "mod", "abs", "fabs" , "gcd", "lcm"]

_trigonometric_functions = [ "sin", "cos", "tan", "arcsin",
                             "arccos", "arctan", "arctan2",
                             "hypot", "sinh", "cosh", "tanh",
                             "arcsinh", "arccosh", "arctanh",
                             "deg2rad", "rad2deg", "degrees",
                             "radians" ]

_bit_twiddling_functions = ["bitwise_and", "bitwise_or",
                            "bitwise_xor", "invert",
                            "left_shift", "right_shift",
                            "bitwise_not" ]

_comparison_functions = [ "greater", "greater_equal", "less",
                          "less_equal", "not_equal", "equal",
                          "logical_and", "logical_or",
                          "logical_xor", "logical_not",
                          "maximum", "minimum", "fmax", "fmin" ]

_floating_functions = [ "isfinite", "isinf", "isnan", "signbit",
                        "copysign", "nextafter", "modf", "ldexp",
                        "frexp", "floor", "ceil", "trunc",
                        "spacing" ]

# This is a set of the ufuncs that are not yet supported by Lowering. In order
# to trigger no-python mode we must not register them until their Lowering is
# implemented.
#
# It also works as a nice TODO list for ufunc support :)
_unsupported = set([ 'frexp', # this one is tricky, as it has 2 returns
                     'modf',  # this one also has 2 returns
                 ])

# A list of ufuncs that are in fact aliases of other ufuncs. They need to insert the
# resolve method, but not register the ufunc itself
_aliases = set(["bitwise_not", "mod", "abs"])

# In python3 np.divide is mapped to np.true_divide
if np.divide == np.true_divide:
    _aliases.add("divide")

def _numpy_ufunc(name):
    func = getattr(np, name)
    class typing_class(Numpy_rules_ufunc):
        key = func

    typing_class.__name__ = "resolve_{0}".format(name)

    if not name in _aliases:
        infer_global(func, types.Function(typing_class))

all_ufuncs = sum([_math_operations, _trigonometric_functions,
                  _bit_twiddling_functions, _comparison_functions,
                  _floating_functions], [])

supported_ufuncs = [x for x in all_ufuncs if x not in _unsupported]

for func in supported_ufuncs:
    _numpy_ufunc(func)

all_ufuncs = [getattr(np, name) for name in all_ufuncs]
supported_ufuncs = [getattr(np, name) for name in supported_ufuncs]

NumpyRulesUnaryArrayOperator.install_operations()
NumpyRulesArrayOperator.install_operations()
NumpyRulesInplaceArrayOperator.install_operations()

supported_array_operators = set(
    NumpyRulesUnaryArrayOperator._op_map.keys()
).union(
    NumpyRulesArrayOperator._op_map.keys()
).union(
    NumpyRulesInplaceArrayOperator._op_map.keys()
)

del _math_operations, _trigonometric_functions, _bit_twiddling_functions
del _comparison_functions, _floating_functions, _unsupported
del _aliases, _numpy_ufunc

# -----------------------------------------------------------------------------
# Install global helpers for array methods.

class Numpy_method_redirection(AbstractTemplate):
    """
    A template redirecting a Numpy global function (e.g. np.sum) to an
    array method of the same name (e.g. ndarray.sum).
    """

    def generic(self, args, kws):
        pysig = None
        if kws:
            if self.method_name == 'sum':
                if 'axis' in kws and 'dtype' not in kws:
                    def sum_stub(arr, axis):
                        pass
                    pysig = utils.pysignature(sum_stub)
                elif 'dtype' in kws and 'axis' not in kws:
                    def sum_stub(arr, dtype):
                        pass
                    pysig = utils.pysignature(sum_stub)
                elif 'dtype' in kws and 'axis' in kws:
                    def sum_stub(arr, axis, dtype):
                        pass
                    pysig = utils.pysignature(sum_stub)
            elif self.method_name == 'argsort':
                def argsort_stub(arr, kind='quicksort'):
                    pass
                pysig = utils.pysignature(argsort_stub)
            else:
                fmt = "numba doesn't support kwarg for {}"
                raise TypingError(fmt.format(self.method_name))

        arr = args[0]
        # This will return a BoundFunction
        meth_ty = self.context.resolve_getattr(arr, self.method_name)
        # Resolve arguments on the bound function
        meth_sig = self.context.resolve_function_type(meth_ty, args[1:], kws)
        if meth_sig is not None:
            return meth_sig.as_function().replace(pysig=pysig)

# Function to glue attributes onto the numpy-esque object
def _numpy_redirect(fname):
    numpy_function = getattr(np, fname)
    cls = type("Numpy_redirect_{0}".format(fname), (Numpy_method_redirection,),
               dict(key=numpy_function, method_name=fname))
    infer_global(numpy_function, types.Function(cls))

for func in ['min', 'max', 'sum', 'prod', 'mean', 'var', 'std',
             'cumsum', 'cumprod', 'argmin', 'argmax', 'argsort',
             'nonzero', 'ravel']:
    _numpy_redirect(func)

# -----------------------------------------------------------------------------
# Numpy scalar constructors

# Register np.int8, etc. as converters to the equivalent Numba types
np_types = set(getattr(np, str(nb_type)) for nb_type in types.number_domain)
np_types.add(np.bool_)
# Those may or may not be aliases (depending on the Numpy build / version)
np_types.add(np.intc)
np_types.add(np.intp)
np_types.add(np.uintc)
np_types.add(np.uintp)

def register_number_classes(register_global):
    for np_type in np_types:
        nb_type = getattr(types, np_type.__name__)

        register_global(np_type, types.NumberClass(nb_type))

register_number_classes(infer_global)

# -----------------------------------------------------------------------------
# Numpy array constructors

def _parse_shape(shape):
    ndim = None
    if isinstance(shape, types.Integer):
        ndim = 1
    elif isinstance(shape, (types.Tuple, types.UniTuple)):
        if all(isinstance(s, types.Integer) for s in shape):
            ndim = len(shape)
    return ndim

def _parse_dtype(dtype):
    if isinstance(dtype, types.DTypeSpec):
        return dtype.dtype

def _parse_nested_sequence(context, typ):
    """
    Parse a (possibly 0d) nested sequence type.
    A (ndim, dtype) tuple is returned.  Note the sequence may still be
    heterogeneous, as long as it converts to the given dtype.
    """
    if isinstance(typ, (types.Buffer,)):
        raise TypingError("%r not allowed in a homogeneous sequence" % typ)
    elif isinstance(typ, (types.Sequence,)):
        n, dtype = _parse_nested_sequence(context, typ.dtype)
        return n + 1, dtype
    elif isinstance(typ, (types.BaseTuple,)):
        if typ.count == 0:
            # Mimick Numpy's behaviour
            return 1, types.float64
        n, dtype = _parse_nested_sequence(context, typ[0])
        dtypes = [dtype]
        for i in range(1, typ.count):
            _n, dtype = _parse_nested_sequence(context, typ[i])
            if _n != n:
                raise TypingError("type %r does not have a regular shape"
                                  % (typ,))
            dtypes.append(dtype)
        dtype = context.unify_types(*dtypes)
        if dtype is None:
            raise TypingError("cannot convert %r to a homogeneous type" % typ)
        return n + 1, dtype
    else:
        # Scalar type => check it's valid as a Numpy array dtype
        as_dtype(typ)
        return 0, typ

@infer_global(np.array)
class NpArray(CallableTemplate):
    """
    Typing template for np.array().
    """

    def generic(self):
        def typer(object, dtype=None):
            ndim, seq_dtype = _parse_nested_sequence(self.context, object)
            if dtype is None:
                dtype = seq_dtype
            else:
                dtype = _parse_dtype(dtype)
                if dtype is None:
                    return
            return types.Array(dtype, ndim, 'C')

        return typer

@infer_global(np.empty)
@infer_global(np.zeros)
@infer_global(np.ones)
class NdConstructor(CallableTemplate):
    """
    Typing template for np.empty(), .zeros(), .ones().
    """

    def generic(self):
        def typer(shape, dtype=None):
            if dtype is None:
                nb_dtype = types.double
            else:
                nb_dtype = _parse_dtype(dtype)

            ndim = _parse_shape(shape)
            if nb_dtype is not None and ndim is not None:
                return types.Array(dtype=nb_dtype, ndim=ndim, layout='C')

        return typer

@infer_global(np.empty_like)
@infer_global(np.zeros_like)
class NdConstructorLike(CallableTemplate):
    """
    Typing template for np.empty_like(), .zeros_like(), .ones_like().
    """

    def generic(self):
        """
        np.empty_like(array) -> empty array of the same shape and layout
        np.empty_like(scalar) -> empty 0-d array of the scalar type
        """
        def typer(arg, dtype=None):
            if dtype is not None:
                nb_dtype = _parse_dtype(dtype)
            elif isinstance(arg, types.Array):
                nb_dtype = arg.dtype
            else:
                nb_dtype = arg
            if nb_dtype is not None:
                if isinstance(arg, types.Array):
                    layout = arg.layout if arg.layout != 'A' else 'C'
                    return arg.copy(dtype=nb_dtype, layout=layout, readonly=False)
                else:
                    return types.Array(nb_dtype, 0, 'C')

        return typer

infer_global(np.ones_like)(NdConstructorLike)

@infer_global(np.full)
class NdFull(CallableTemplate):

    def generic(self):
        def typer(shape, fill_value, dtype=None):
            if dtype is None:
                nb_dtype = fill_value
            else:
                nb_dtype = _parse_dtype(dtype)

            ndim = _parse_shape(shape)
            if nb_dtype is not None and ndim is not None:
                return types.Array(dtype=nb_dtype, ndim=ndim, layout='C')

        return typer

@infer_global(np.full_like)
class NdFullLike(CallableTemplate):

    def generic(self):
        """
        np.full_like(array, val) -> array of the same shape and layout
        np.full_like(scalar, val) -> 0-d array of the scalar type
        """
        def typer(arg, fill_value, dtype=None):
            if dtype is not None:
                nb_dtype = _parse_dtype(dtype)
            elif isinstance(arg, types.Array):
                nb_dtype = arg.dtype
            else:
                nb_dtype = arg
            if nb_dtype is not None:
                if isinstance(arg, types.Array):
                    return arg.copy(dtype=nb_dtype, readonly=False)
                else:
                    return types.Array(dtype=nb_dtype, ndim=0, layout='C')

        return typer

@infer_global(np.identity)
class NdIdentity(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws
        n = args[0]
        if not isinstance(n, types.Integer):
            return
        if len(args) >= 2:
            nb_dtype = _parse_dtype(args[1])
        else:
            nb_dtype = types.float64

        if nb_dtype is not None:
            return_type = types.Array(ndim=2, dtype=nb_dtype, layout='C')
            return signature(return_type, *args)

def _infer_dtype_from_inputs(inputs):
    return dtype

@infer_global(np.linspace)
class NdLinspace(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws
        bounds = args[:2]
        if not all(isinstance(arg, types.Number) for arg in bounds):
            return
        if len(args) >= 3:
            num = args[2]
            if not isinstance(num, types.Integer):
                return
        if len(args) >= 4:
            # Not supporting the other arguments as it would require
            # keyword arguments for reasonable use.
            return
        if any(isinstance(arg, types.Complex) for arg in bounds):
            dtype = types.complex128
        else:
            dtype = types.float64
        return_type = types.Array(ndim=1, dtype=dtype, layout='C')
        return signature(return_type, *args)

@infer_global(np.frombuffer)
class NdFromBuffer(CallableTemplate):

    def generic(self):
        def typer(buffer, dtype=None):
            if not isinstance(buffer, types.Buffer) or buffer.layout != 'C':
                return
            if dtype is None:
                nb_dtype = types.float64
            else:
                nb_dtype = _parse_dtype(dtype)

            if nb_dtype is not None:
                return types.Array(dtype=nb_dtype, ndim=1, layout='C',
                                   readonly=not buffer.mutable)

        return typer

@infer_global(np.sort)
class NdSort(CallableTemplate):

    def generic(self):
        def typer(a):
            if isinstance(a, types.Array) and a.ndim == 1:
                return a

        return typer

@infer_global(np.asfortranarray)
class AsFortranArray(CallableTemplate):

    def generic(self):
        def typer(a):
            if isinstance(a, types.Array):
                return a.copy(layout='F', ndim=max(a.ndim, 1))

        return typer

@infer_global(np.ascontiguousarray)
class AsContiguousArray(CallableTemplate):

    def generic(self):
        def typer(a):
            if isinstance(a, types.Array):
                return a.copy(layout='C', ndim=max(a.ndim, 1))

        return typer

@infer_global(np.copy)
class NdCopy(CallableTemplate):

    def generic(self):
        def typer(a):
            if isinstance(a, types.Array):
                layout = 'F' if a.layout == 'F' else 'C'
                return a.copy(layout=layout, readonly=False)

        return typer

@infer_global(np.expand_dims)
class NdExpandDims(CallableTemplate):

    def generic(self):
        def typer(a, axis):
            if (not isinstance(a, types.Array)
                or not isinstance(axis, types.Integer)):
                return

            layout = a.layout if a.ndim <= 1 else 'A'
            return a.copy(ndim=a.ndim + 1, layout=layout)

        return typer

class BaseAtLeastNdTemplate(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws
        if not args or not all(isinstance(a, types.Array) for a in args):
            return

        rets = [self.convert_array(a) for a in args]
        if len(rets) > 1:
            retty = types.BaseTuple.from_types(rets)
        else:
            retty = rets[0]
        return signature(retty, *args)

@infer_global(np.atleast_1d)
class NdAtLeast1d(BaseAtLeastNdTemplate):

    def convert_array(self, a):
        return a.copy(ndim=max(a.ndim, 1))

@infer_global(np.atleast_2d)
class NdAtLeast2d(BaseAtLeastNdTemplate):

    def convert_array(self, a):
        return a.copy(ndim=max(a.ndim, 2))

@infer_global(np.atleast_3d)
class NdAtLeast3d(BaseAtLeastNdTemplate):

    def convert_array(self, a):
        return a.copy(ndim=max(a.ndim, 3))

def _homogeneous_dims(context, func_name, arrays):
    ndim = arrays[0].ndim
    for a in arrays:
        if a.ndim != ndim:
            raise TypeError("%s(): all the input arrays "
                            "must have same number of dimensions"
                            % func_name)
    return ndim

def _sequence_of_arrays(context, func_name, arrays,
                        dim_chooser=_homogeneous_dims):
    if (not isinstance(arrays, types.BaseTuple)
        or not len(arrays)
        or not all(isinstance(a, types.Array) for a in arrays)):
        raise TypeError("%s(): expecting a non-empty tuple of arrays, "
                        "got %s" % (func_name, arrays))

    ndim = dim_chooser(context, func_name, arrays)

    dtype = context.unify_types(*(a.dtype for a in arrays))
    if dtype is None:
        raise TypeError("%s(): input arrays must have "
                        "compatible dtypes" % func_name)

    return dtype, ndim

def _choose_concatenation_layout(arrays):
    # Only create a F array if all input arrays have F layout.
    # This is a simplified version of Numpy's behaviour,
    # while Numpy's actually processes the input strides to
    # decide on optimal output strides
    # (see PyArray_CreateMultiSortedStridePerm()).
    return 'F' if all(a.layout == 'F' for a in arrays) else 'C'

@infer_global(np.concatenate)
class NdConcatenate(CallableTemplate):

    def generic(self):
        def typer(arrays, axis=None):
            if axis is not None and not isinstance(axis, types.Integer):
                # Note Numpy allows axis=None, but it isn't documented:
                # https://github.com/numpy/numpy/issues/7968
                return

            dtype, ndim = _sequence_of_arrays(self.context,
                                              "np.concatenate", arrays)
            if ndim == 0:
                raise TypeError("zero-dimensional arrays cannot be concatenated")

            layout = _choose_concatenation_layout(arrays)

            return types.Array(dtype, ndim, layout)

        return typer

@infer_global(np.stack)
class NdStack(CallableTemplate):

    def generic(self):
        def typer(arrays, axis=None):
            if axis is not None and not isinstance(axis, types.Integer):
                # Note Numpy allows axis=None, but it isn't documented:
                # https://github.com/numpy/numpy/issues/7968
                return

            dtype, ndim = _sequence_of_arrays(self.context,
                                                "np.stack", arrays)

            # This diverges from Numpy's behaviour, which simply inserts
            # a new stride at the requested axis (therefore can return
            # a 'A' array).
            layout = 'F' if all(a.layout == 'F' for a in arrays) else 'C'

            return types.Array(dtype, ndim + 1, layout)

        return typer

class BaseStackTemplate(CallableTemplate):

    def generic(self):
        def typer(arrays):
            dtype, ndim = _sequence_of_arrays(self.context,
                                              self.func_name, arrays)

            ndim = max(ndim, self.ndim_min)
            layout = _choose_concatenation_layout(arrays)

            return types.Array(dtype, ndim, layout)

        return typer

@infer_global(np.hstack)
class NdStack(BaseStackTemplate):
    func_name = "np.hstack"
    ndim_min = 1

@infer_global(np.vstack)
class NdStack(BaseStackTemplate):
    func_name = "np.vstack"
    ndim_min = 2

@infer_global(np.dstack)
class NdStack(BaseStackTemplate):
    func_name = "np.dstack"
    ndim_min = 3

def _column_stack_dims(context, func_name, arrays):
    # column_stack() allows stacking 1-d and 2-d arrays together
    for a in arrays:
        if a.ndim < 1 or a.ndim > 2:
            raise TypeError("np.column_stack() is only defined on "
                            "1-d and 2-d arrays")
    return 2

@infer_global(np.column_stack)
class NdColumnStack(CallableTemplate):

    def generic(self):
        def typer(arrays):
            dtype, ndim = _sequence_of_arrays(self.context,
                                              "np.column_stack", arrays,
                                              dim_chooser=_column_stack_dims)

            layout = _choose_concatenation_layout(arrays)

            return types.Array(dtype, ndim, layout)

        return typer

# -----------------------------------------------------------------------------
# Linear algebra

class MatMulTyperMixin(object):

    def matmul_typer(self, a, b, out=None):
        """
        Typer function for Numpy matrix multiplication.
        """
        if not isinstance(a, types.Array) or not isinstance(b, types.Array):
            return
        if not all(x.ndim in (1, 2) for x in (a, b)):
            raise TypingError("%s only supported on 1-D and 2-D arrays"
                              % (self.func_name, ))
        # Output dimensionality
        ndims = set([a.ndim, b.ndim])
        if ndims == set([2]):
            # M * M
            out_ndim = 2
        elif ndims == set([1, 2]):
            # M* V and V * M
            out_ndim = 1
        elif ndims == set([1]):
            # V * V
            out_ndim = 0

        if out is not None:
            if out_ndim == 0:
                raise TypeError("explicit output unsupported for vector * vector")
            elif out.ndim != out_ndim:
                raise TypeError("explicit output has incorrect dimensionality")
            if not isinstance(out, types.Array) or out.layout != 'C':
                raise TypeError("output must be a C-contiguous array")
            all_args = (a, b, out)
        else:
            all_args = (a, b)

        if not (config.DISABLE_PERFORMANCE_WARNINGS or
                all(x.layout in 'CF' for x in (a, b))):
            msg = ("%s is faster on contiguous arrays, called on %s" %
                   (self.func_name, (a, b)))
            warnings.warn(NumbaPerformanceWarning(msg))
        if not all(x.dtype == a.dtype for x in all_args):
            raise TypingError("%s arguments must all have "
                              "the same dtype" % (self.func_name,))
        if not isinstance(a.dtype, (types.Float, types.Complex)):
            raise TypingError("%s only supported on "
                              "float and complex arrays"
                              % (self.func_name,))
        if out:
            return out
        elif out_ndim > 0:
            return types.Array(a.dtype, out_ndim, 'C')
        else:
            return a.dtype

@infer_global(np.dot)
class Dot(MatMulTyperMixin, CallableTemplate):
    func_name = "np.dot()"

    def generic(self):
        def typer(a, b, out=None):
            # NOTE: np.dot() and the '@' operator have distinct semantics
            # for >2-D arrays, but we don't support them.
            return self.matmul_typer(a, b, out)

        return typer

@infer_global(np.vdot)
class VDot(CallableTemplate):

    def generic(self):
        def typer(a, b):
            if not isinstance(a, types.Array) or not isinstance(b, types.Array):
                return
            if not all(x.ndim == 1 for x in (a, b)):
                raise TypingError("np.vdot() only supported on 1-D arrays")
            if not all(x.layout in 'CF' for x in (a, b)):
                warnings.warn("np.vdot() is faster on contiguous arrays, called on %s"
                              % ((a, b),), NumbaPerformanceWarning)
            if not all(x.dtype == a.dtype for x in (a, b)):
                raise TypingError("np.vdot() arguments must all have "
                                  "the same dtype")
            if not isinstance(a.dtype, (types.Float, types.Complex)):
                raise TypingError("np.vdot() only supported on "
                                  "float and complex arrays")
            return a.dtype

        return typer

@infer_global(operator.matmul)
class MatMul(MatMulTyperMixin, AbstractTemplate):
    key = operator.matmul
    func_name = "'@'"

    def generic(self, args, kws):
        assert not kws
        restype = self.matmul_typer(*args)
        if restype is not None:
            return signature(restype, *args)

def _check_linalg_matrix(a, func_name):
    if not isinstance(a, types.Array):
        return
    if not a.ndim == 2:
        raise TypingError("np.linalg.%s() only supported on 2-D arrays"
                          % func_name)
    if not isinstance(a.dtype, (types.Float, types.Complex)):
        raise TypingError("np.linalg.%s() only supported on "
                          "float and complex arrays" % func_name)

# -----------------------------------------------------------------------------
# Miscellaneous functions

@infer_global(np.ndenumerate)
class NdEnumerate(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws
        arr, = args

        if isinstance(arr, types.Array):
            enumerate_type = types.NumpyNdEnumerateType(arr)
            return signature(enumerate_type, *args)

@infer_global(np.nditer)
class NdIter(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws
        if len(args) != 1:
            return
        arrays, = args

        if isinstance(arrays, types.BaseTuple):
            if not arrays:
                return
            arrays = list(arrays)
        else:
            arrays = [arrays]
        nditerty = types.NumpyNdIterType(arrays)
        return signature(nditerty, *args)

@infer_global(pndindex)
@infer_global(np.ndindex)
class NdIndex(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws

        # Either ndindex(shape) or ndindex(*shape)
        if len(args) == 1 and isinstance(args[0], types.BaseTuple):
            tup = args[0]
            if tup.count > 0 and not isinstance(tup, types.UniTuple):
                # Heterogeneous tuple
                return
            shape = list(tup)
        else:
            shape = args

        if all(isinstance(x, types.Integer) for x in shape):
            iterator_type = types.NumpyNdIndexType(len(shape))
            return signature(iterator_type, *args)

# We use the same typing key for np.round() and np.around() to
# re-use the implementations automatically.
@infer_global(np.round)
@infer_global(np.around, typing_key=np.round)
class Round(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws
        assert 1 <= len(args) <= 3

        arg = args[0]
        if len(args) == 1:
            decimals = types.intp
            out = None
        else:
            decimals = args[1]
            if len(args) == 2:
                out = None
            else:
                out = args[2]

        supported_scalars = (types.Integer, types.Float, types.Complex)
        if isinstance(arg, supported_scalars):
            assert out is None
            return signature(arg, *args)
        if (isinstance(arg, types.Array) and isinstance(arg.dtype, supported_scalars) and
            isinstance(out, types.Array) and isinstance(out.dtype, supported_scalars) and
            out.ndim == arg.ndim):
            # arg can only be complex if out is complex too
            if (not isinstance(arg.dtype, types.Complex)
                or isinstance(out.dtype, types.Complex)):
                return signature(out, *args)

@infer_global(np.where)
class Where(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws

        if len(args) == 1:
            # 0-dim arrays return one result array
            ary = args[0]
            ndim = max(ary.ndim, 1)
            retty = types.UniTuple(types.Array(types.intp, 1, 'C'), ndim)
            return signature(retty, ary)

        elif len(args) == 3:
            cond, x, y = args
            retdty = from_dtype(np.promote_types(
                        as_dtype(getattr(args[1], 'dtype', args[1])),
                        as_dtype(getattr(args[2], 'dtype', args[2]))))
            if isinstance(cond, types.Array):
                # array where()
                if isinstance(x, types.Array) and isinstance(y, types.Array):
                    if (cond.ndim == x.ndim == y.ndim):
                        if x.layout == y.layout == cond.layout:
                            retty = types.Array(retdty, x.ndim, x.layout)
                        else:
                            retty = types.Array(retdty, x.ndim, 'C')
                        return signature(retty, *args)
                else:
                    # x and y both scalar
                    retty = types.Array(retdty, cond.ndim, cond.layout)
                    return signature(retty, *args)
            else:
                # scalar where()
                if not isinstance(x, types.Array):
                    retty = types.Array(retdty, 0, 'C')
                    return signature(retty, *args)

@infer_global(np.sinc)
class Sinc(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws
        assert len(args) == 1
        arg = args[0]
        supported_scalars = (types.Float, types.Complex)
        if (isinstance(arg, supported_scalars) or
              (isinstance(arg, types.Array) and
               isinstance(arg.dtype, supported_scalars))):
            return signature(arg, arg)

@infer_global(np.angle)
class Angle(CallableTemplate):
    """
    Typing template for np.angle()
    """
    def generic(self):
        def typer(z, deg=False):
            if isinstance(z, types.Array):
                dtype = z.dtype
            else:
                dtype = z
            if isinstance(dtype, types.Complex):
                ret_dtype = dtype.underlying_float
            elif isinstance(dtype, types.Float):
                ret_dtype = dtype
            else:
                return
            if isinstance(z, types.Array):
                return z.copy(dtype=ret_dtype)
            else:
                return ret_dtype
        return typer

@infer_global(np.diag)
class DiagCtor(CallableTemplate):
    """
    Typing template for np.diag()
    """
    def generic(self):
        def typer(ref, k=0):
            if isinstance(ref, types.Array):
                if ref.ndim == 1:
                    rdim = 2
                elif ref.ndim == 2:
                    rdim = 1
                else:
                    return None
                if isinstance(k, (int, types.Integer)):
                    return types.Array(ndim=rdim, dtype=ref.dtype, layout='C')
        return typer

@infer_global(np.take)
class Take(AbstractTemplate):

    def generic(self, args, kws):
        assert not kws
        assert len(args) == 2
        arr, ind = args
        if isinstance(ind, types.Number):
            retty = arr.dtype
        elif isinstance(ind, types.Array):
            retty = types.Array(ndim=ind.ndim, dtype=arr.dtype, layout='C')
        elif isinstance(ind, types.List):
            retty = types.Array(ndim=1, dtype=arr.dtype, layout='C')
        elif isinstance(ind, types.BaseTuple):
            retty = types.Array(ndim=np.ndim(ind), dtype=arr.dtype, layout='C')
        else:
            return None

        return signature(retty, *args)

# -----------------------------------------------------------------------------
# Numba helpers

@infer_global(carray)
class NumbaCArray(CallableTemplate):
    layout = 'C'

    def generic(self):
        func_name = self.key.__name__

        def typer(ptr, shape, dtype=types.none):
            if ptr is types.voidptr:
                ptr_dtype = None
            elif isinstance(ptr, types.CPointer):
                ptr_dtype = ptr.dtype
            else:
                raise TypeError("%s(): pointer argument expected, got '%s'"
                                % (func_name, ptr))

            if dtype is types.none:
                if ptr_dtype is None:
                    raise TypeError("%s(): explicit dtype required for void* argument"
                                    % (func_name,))
                dtype = ptr_dtype
            elif isinstance(dtype, types.DTypeSpec):
                dtype = dtype.dtype
                if ptr_dtype is not None and dtype != ptr_dtype:
                    raise TypeError("%s(): mismatching dtype '%s' for pointer type '%s'"
                                    % (func_name, dtype, ptr))
            else:
                raise TypeError("%s(): invalid dtype spec '%s'"
                                % (func_name, dtype))

            ndim = _parse_shape(shape)
            if ndim is None:
                raise TypeError("%s(): invalid shape '%s'"
                                % (func_name, shape))

            return types.Array(dtype, ndim, self.layout)

        return typer

@infer_global(farray)
class NumbaFArray(NumbaCArray):
    layout = 'F'


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Implementation for Single Image Haze Removal Using Dark Channel Prior.

Reference:
http://research.microsoft.com/en-us/um/people/kahe/cvpr09/
http://research.microsoft.com/en-us/um/people/kahe/eccv10/
"""

import numpy as np
from PIL import Image

from guidedfilter import guided_filter

R, G, B = 0, 1, 2  # index for convenience
L = 256  # color depth

def get_dark_channel(I, w):
    """Get the dark channel prior in the (RGB) image data.

    Parameters
    -----------
    I:  an M * N * 3 numpy array containing data ([0, L-1]) in the image where
        M is the height, N is the width, 3 represents R/G/B channels.
    w:  window size

    Return
    -----------
    An M * N array for the dark channel prior ([0, L-1]).
    """
    M, N, _ = I.shape
    padded = np.pad(I, ((w / 2, w / 2), (w / 2, w / 2), (0, 0)), 'edge')
    darkch = np.zeros((M, N))
    for i, j in np.ndindex(darkch.shape):
        darkch[i, j] = np.min(padded[i:i + w, j:j + w, :])  # CVPR09, eq.5
    return darkch

def get_atmosphere(I, darkch, p):
    """Get the atmosphere light in the (RGB) image data.

    Parameters
    -----------
    I:      the M * N * 3 RGB image data ([0, L-1]) as numpy array
    darkch: the dark channel prior of the image as an M * N numpy array
    p:      percentage of pixels for estimating the atmosphere light

    Return
    -----------
    A 3-element array containing atmosphere light ([0, L-1]) for each channel
    """
    # reference CVPR09, 4.4
    M, N = darkch.shape
    flatI = I.reshape(M * N, 3)
    flatdark = darkch.ravel()
    searchidx = (-flatdark).argsort()[:M * N * p]  # find top M * N * p indexes
    print('atmosphere light region:', [(i / N, i % N) for i in searchidx])

    # return the highest intensity for each channel
    return np.max(flatI.take(searchidx, axis=0), axis=0)

def get_transmission(I, A, darkch, omega, w):
    """Get the transmission esitmate in the (RGB) image data.

    Parameters
    -----------
    I:       the M * N * 3 RGB image data ([0, L-1]) as numpy array
    A:       a 3-element array containing atmosphere light
             ([0, L-1]) for each channel
    darkch:  the dark channel prior of the image as an M * N numpy array
    omega:   bias for the estimate
    w:       window size for the estimate

    Return
    -----------
    An M * N array containing the transmission rate ([0.0, 1.0])
    """
    return 1 - omega * get_dark_channel(I / A, w)  # CVPR09, eq.12

def dehaze_raw(I, tmin=0.2, Amax=220, w=15, p=0.0001,
               omega=0.95, guided=True, r=40, eps=1e-3):
    """Get the dark channel prior, atmosphere light, transmission rate
       and refined transmission rate for raw RGB image data.

    Parameters
    -----------
    I:      M * N * 3 data as numpy array for the hazy image
    tmin:   threshold of transmission rate
    Amax:   threshold of atmosphere light
    w:      window size of the dark channel prior
    p:      percentage of pixels for estimating the atmosphere light
    omega:  bias for the transmission estimate

    guided: whether to use the guided filter to fine the image
    r:      the radius of the guidance
    eps:    epsilon for the guided filter

    Return
    -----------
    (Idark, A, rawt, refinedt) if guided=False, then rawt == refinedt
    """
    m, n, _ = I.shape
    Idark = get_dark_channel(I, w)

    A = get_atmosphere(I, Idark, p)
    A = np.minimum(A, Amax)  # threshold A
    print('atmosphere', A)

    rawt = get_transmission(I, A, Idark, omega, w)
    print('raw transmission rate',)
    print('between [%.4f, %.4f]' % (rawt.min(), rawt.max()))

    rawt = refinedt = np.maximum(rawt, tmin)  # threshold t
    if guided:
        normI = (I - I.min()) / (I.max() - I.min())  # normalize I
        refinedt = guided_filter(normI, refinedt, r, eps)

    print('refined transmission rate',)
    print('between [%.4f, %.4f]' % (refinedt.min(), refinedt.max()))

    return Idark, A, rawt, refinedt

def get_radiance(I, A, t):
    """Recover the radiance from raw image data with atmosphere light
       and transmission rate estimate.

    Parameters
    ----------
    I:      M * N * 3 data as numpy array for the hazy image
    A:      a 3-element array containing atmosphere light
            ([0, L-1]) for each channel
    t:      estimate fothe transmission rate

    Return
    ----------
    M * N * 3 numpy array for the recovered radiance
    """
    tiledt = np.zeros_like(I)  # tiled to M * N * 3
    tiledt[:, :, R] = tiledt[:, :, G] = tiledt[:, :, B] = t
    return (I - A) / tiledt + A  # CVPR09, eq.16

def dehaze(im, tmin=0.2, Amax=220, w=15, p=0.0001,
           omega=0.95, guided=True, r=40, eps=1e-3):
    """Dehaze the given RGB image.

    Parameters
    ----------
    im:     the Image object of the RGB image
    guided: refine the dehazing with guided filter or not
    other parameters are the same as `dehaze_raw`

    Return
    ----------
    (dark, rawt, refinedt, rawrad, rerad)
    Images for dark channel prior, raw transmission estimate,
    refiend transmission estimate, recovered radiance with raw t,
    recovered radiance with refined t.
    """
    I = np.asarray(im, dtype=np.float64)
    Idark, A, rawt, refinedt = dehaze_raw(I, tmin, Amax, w, p,
                                          omega, guided, r, eps)
    white = np.full_like(Idark, L - 1)

    def to_img(raw):
        # threshold to [0, L-1]
        cut = np.maximum(np.minimum(raw, L - 1), 0).astype(np.uint8)

        if len(raw.shape) == 3:
            print('Range for each channel:')
            for ch in range(3):
                print('[%.2f, %.2f]' % (raw[:, :, ch].max(), raw[:, :, ch].min()))
            return Image.fromarray(cut)
        else:
            return Image.fromarray(cut)

    return [to_img(raw) for raw in (Idark, white * rawt, white * refinedt,
                                    get_radiance(I, A, rawt),
                                    get_radiance(I, A, refinedt))]


"""

Source.py

Author: Jordan Mirocha
Affiliation: University of Colorado at Boulder
Created on: Sun Jul 22 16:28:08 2012

Description: Initialize a radiation source.

"""

import re, os
import numpy as np
from scipy.integrate import quad
from ..util import ParameterFile
from ..physics.Hydrogen import Hydrogen
from ..physics.Cosmology import Cosmology
from ..physics.Constants import erg_per_ev, E_LL
from ..static.IntegralTables import IntegralTable
from ..static.InterpolationTables import LookupTable
from ..util.SetDefaultParameterValues import SourceParameters, \
    CosmologyParameters
from ..physics.CrossSections import PhotoIonizationCrossSection as sigma_E

try:
    import h5py
except ImportError:
    pass

np.seterr(all='ignore')   # exp overflow occurs when integrating BB
                          # will return 0 as it should for x large

cosmo_pars = CosmologyParameters()

class Source(object):
    def __init__(self, grid=None, logN=None, init_tabs=True):
        """
        Initialize a radiation source object.

        ..note:: This is inherited by all other ares.sources classes.

        Parameters
        ----------
        grid: rt1d.static.Grid.Grid instance
        logN: column densities over which to tabulate integral quantities

        """

        # Update cosmological parameters
        # Why is this necessary? J.M.12.27.2015
        for par in cosmo_pars:
            if par in self.pf:
                continue

            self.pf[par] = cosmo_pars[par]

        # Modify parameter file if spectrum_file provided
        #self._load_spectrum()

        # Correct emission limits if none were provided
        self.Emin = self.pf['source_Emin']
        self.Emax = self.pf['source_Emax']
        self.logEmin = np.log10(self.Emin)
        self.logEmax = np.log10(self.Emax)

        if self.pf['source_EminNorm'] == None:
            self.pf['source_EminNorm'] = self.pf['source_Emin']
        if self.pf['source_EmaxNorm'] == None:
            self.pf['source_EmaxNorm'] = self.pf['source_Emax']

        self.EminNorm = self.pf['source_EminNorm']
        self.EmaxNorm = self.pf['source_EmaxNorm']

        # Number of frequencies
        #if self.discrete:
        #    self.E = np.array(self.pf['source_E'])
        #    self.LE = np.array(self.pf['source_LE'])
        #    self.Nfreq = len(self.E)
        #
        #if self.src._name == 'DiffuseSource':
        #    self.ionization_rate = self.src.ionization_rate
        #    self.secondary_ionization_rate = self.src.secondary_ionization_rate
        #    self.heating_rate = self.src.heating_rate
        #
        #self.Lbol = self.Lbol0 = self.BolometricLuminosity(0.0)

        # Create lookup tables for integral quantities
        if init_tabs and grid is not None:
            self._create_integral_table(logN=logN)

    @property
    def cosm(self):
        if not hasattr(self, '_cosm'):
            if self.grid is None:
                self._cosm = Cosmology(
                    omega_m_0=self.pf['omega_m_0'],
                    omega_l_0=self.pf['omega_l_0'],
                    omega_b_0=self.pf['omega_b_0'],
                    hubble_0=self.pf['hubble_0'],
                    helium_by_number=self.pf['helium_by_number'],
                    cmb_temp_0=self.pf['cmb_temp_0'],
                    approx_highz=self.pf['approx_highz'],
                    sigma_8=self.pf['sigma_8'],
                    primordial_index=self.pf['primordial_index'])
            else:
                self._cosm = self.grid.cosm

        return self._cosm

    @property
    def multi_freq(self):
        if not hasattr(self, '_multi_freq'):
            self._multi_freq = self.discrete and not self.pf['source_multigroup']

        return self._multi_freq

    @property
    def multi_group(self):
        if not hasattr(self, '_multi_group'):
            self._multi_group = self.discrete and self.pf['source_multigroup']

        return self._multi_group

    @property
    def ionizing(self):
        # See if source emits ionizing photons
        # Should also be function of absorbers
        if not hasattr(self, '_ionizing'):
            self._ionizing = self.pf['source_Emax'] > E_LL

        return self._ionizing

    @property
    def grid(self):
        if not hasattr(self, '_grid'):
            self._grid = None

        return self._grid

    @grid.setter
    def grid(self, value):
        self._grid = value

    @property
    def discrete(self):
        if not hasattr(self, '_discrete'):
            self._discrete = (self.pf['source_E'] != None) #\
                  #or self.pf['optically_thin']

        return self._discrete

    @property
    def continuous(self):
        if not hasattr(self, '_continuous'):
            self._continuous = not self.discrete

        return self._continuous

    @property
    def hydr(self):
        if not hasattr(self, '_hydr'):
            self._hydr = None

        return self._hydr

    @hydr.setter
    def hydr(self, value):
        self._hydr = value

    @property
    def frec(self):
        """
        Compute average recycling fraction (i.e., spectrum-weighted frec).
        """

        if self.hydr is None:
            return None

        n = np.arange(2, self.hydr.nmax)
        En = np.array(map(self.hydr.ELyn, n))
        In = np.array(map(self.Spectrum, En)) / En
        fr = np.array(map(self.hydr.frec, n))

        return np.sum(fr * In) / np.sum(In)

    @property
    def intrinsic_hardening(self):
        if not hasattr(self, '_intrinsic_hardening'):
            if 'source_hardening' in self.pf:
                self._intrinsic_hardening = \
                    self.pf['source_hardening'] == 'intrinsic'
            else:
                self._intrinsic_hardening = False

        return self._intrinsic_hardening

    def _hardening_factor(self, E):
        return np.exp(-10.**self.logN \
            * (sigma_E(E, 0) + self.cosm.y * sigma_E(E, 1)))

    @property
    def logN(self):
        if not hasattr(self, '_logN'):
            if 'source_logN' in self.pf:
                self._logN = self.pf['source_logN']
            else:
                self._logN = -np.inf

        return self._logN

    @property
    def _normL(self):
        if not hasattr(self, '_normL_'):
            if self.intrinsic_hardening:
                self._normL_ = 1. / quad(self._Intensity,
                    self.pf['source_EminNorm'], self.pf['source_EmaxNorm'])[0]
            else:
                integrand = lambda EE: self._Intensity(EE) / self._hardening_factor(EE)
                self._normL_ = 1. / quad(integrand,
                    self.pf['source_EminNorm'], self.pf['source_EmaxNorm'])[0]

        return self._normL_

    #def _load_spectrum(self):
    #    """ Modify a few parameters if spectrum_file provided. """
    #
    #    fn = self.pf['spectrum_file']
    #
    #    if fn is None:
    #        return
    #
    #    # Read spectrum - expect hdf5 with (at least) E, LE, and t datasets.
    #    if re.search('.hdf5', fn):
    #        f = h5py.File(fn)
    #        try:
    #            self.pf['tables_times'] = f['t'].value
    #        except:
    #            self.pf['tables_times'] = None
    #            self.pf['spectrum_evolving'] = False
    #
    #        self.pf['spectrum_E'] = f['E'].value
    #        self.pf['spectrum_LE'] = f['LE'].value
    #        f.close()
    #
    #        if len(self.pf['spectrum_LE'].shape) > 1 \
    #            and not self.pf['spectrum_evolving']:
    #            self.pf['spectrum_LE'] = self.pf['spectrum_LE'][0]
    #    else:
    #        spec = readtab(fn)
    #        if len(spec) == 2:
    #            self.pf['spectrum_E'], self.pf['spectrum_LE'] = spec
    #        else:
    #            self.pf['spectrum_E'], self.pf['spectrum_LE'], \
    #                self.pf['spectrum_t'] = spec

    @property
    def tables(self):
        if not hasattr(self, '_tables'):
            self._create_integral_table()
        return self._tables

    @property
    def tab(self):
        if not hasattr(self, '_tab'):
            self._create_integral_table()
        return self._tab

    @property
    def tabs(self):
        if not hasattr(self, '_tabs'):
            self._create_integral_table()
        return self._tabs

    def _create_integral_table(self, logN=None):
        """
        Take tables and create interpolation functions.
        """

        if self.discrete:
            return

        if self._name == 'diffuse':
            return

        if self.pf['source_table'] is None:
            # Overide defaults if supplied - this is dangerous
            if logN is not None:
                self.pf.update({'tables_dlogN': [np.diff(tmp) for tmp in logN]})
                self.pf.update({'tables_logNmin': [np.min(tmp) for tmp in logN]})
                self.pf.update({'tables_logNmax': [np.max(tmp) for tmp in logN]})

            # Tabulate away!
            self._tab = IntegralTable(self.pf, self, self.grid, logN)
            self._tabs = self.tab.TabulateRateIntegrals()
        else:
            self._tab = IntegralTable(self.pf, self, self.grid, logN)
            self._tabs = self.tab.load(self.pf['source_table'])

        self._setup_interp()

    def _setup_interp(self):
        self._tables = {}
        for tab in self.tabs:
            self._tables[tab] = \
                LookupTable(self.pf, tab, self.tab.logN, self.tabs[tab],
                    self.tab.logx, self.tab.t)

    @property
    def sigma(self):
        """
        Compute bound-free absorption cross-section for all frequencies.
        """
        if not self.discrete:
            return None
        if not hasattr(self, '_sigma_all'):
            self._sigma_all = np.array(map(sigma_E, self.E))

        return self._sigma_all

    @property
    def Qdot(self):
        """
        Returns number of photons emitted (s^-1) at all frequencies.
        """
        if not hasattr(self, '_Qdot_all'):
            self._Qdot_all = self.Lbol * self.LE / self.E / erg_per_ev

        return self._Qdot_all

    @property
    def hnu_bar(self):
        """
        Average ionizing (per absorber) photon energy in eV.
        """
        if not hasattr(self, '_hnu_bar_all'):
            self._hnu_bar_all = np.zeros_like(self.grid.zeros_absorbers)
            self._qdot_bar_all = np.zeros_like(self.grid.zeros_absorbers)
            for i, absorber in enumerate(self.grid.absorbers):
                self._hnu_bar_all[i], self._qdot_bar_all[i] = \
                    self._FrequencyAveragedBin(absorber=absorber)

        return self._hnu_bar_all

    def AveragePhotonEnergy(self, Emin, Emax):
        """
        Return average photon energy in supplied band.
        """

        integrand = lambda EE: self.Spectrum(EE) * EE
        norm = lambda EE: self.Spectrum(EE)

        return quad(integrand, Emin, Emax)[0] / quad(norm, Emin, Emax)[0]

    @property
    def qdot_bar(self):
        """
        Average ionizing photon luminosity (per absorber) in s^-1.
        """
        if not hasattr(self, '_qdot_bar_all'):
            hnu_bar = self.hnu_bar

        return self._qdot_bar_all

    @property
    def sigma_bar(self):
        """
        Frequency averaged cross section (single bandpass).
        """
        if not hasattr(self, '_sigma_bar_all'):
            self._sigma_bar_all = np.zeros_like(self.grid.zeros_absorbers)
            for i, absorber in enumerate(self.grid.absorbers):
                integrand = lambda x: self.Spectrum(x) \
                    * self.grid.bf_cross_sections[absorber](x) / x

                self._sigma_bar_all[i] = self.Lbol \
                    * quad(integrand, self.grid.ioniz_thresholds[absorber],
                      self.Emax)[0] / self.qdot_bar[i] / erg_per_ev

        return self._sigma_bar_all

    @property
    def sigma_tilde(self):
        if not hasattr(self, '_sigma_tilde_all'):
            self._sigma_tilde_all = np.zeros_like(self.grid.zeros_absorbers)
            for i, absorber in enumerate(self.grid.absorbers):
                integrand = lambda x: self.Spectrum(x) \
                    * self.grid.bf_cross_sections[absorber](x)
                self._sigma_tilde_all[i] = quad(integrand,
                    self.grid.ioniz_thresholds[absorber], self.Emax)[0] \
                    / self.fLbol_ionizing[i]

        return self._sigma_tilde_all

    @property
    def fLbol_ionizing(self):
        """
        Fraction of bolometric luminosity emitted above all ionization
        thresholds.
        """
        if not hasattr(self, '_fLbol_ioniz_all'):
            self._fLbol_ioniz_all = np.zeros_like(self.grid.zeros_absorbers)
            for i, absorber in enumerate(self.grid.absorbers):
                self._fLbol_ioniz_all[i] = quad(self.Spectrum,
                    self.grid.ioniz_thresholds[absorber], self.Emax)[0]

        return self._fLbol_ioniz_all

    @property
    def Gamma_bar(self):
        """
        Return ionization rate (as a function of radius) assuming optical
        depth to cells and of cells is small.
        """
        if not hasattr(self, '_Gamma_bar_all'):
            self._Gamma_bar_all = \
                np.zeros([self.grid.dims, self.grid.N_absorbers])
            for i, absorber in enumerate(self.grid.absorbers):
                self._Gamma_bar_all[..., i] = self.Lbol * self.sigma_bar[i] \
                    * self.fLbol_ionizing[i] / 4. / np.pi / self.grid.r_mid**2 \
                    / self.hnu_bar[i] / erg_per_ev

        return self._Gamma_bar_all

    @property
    def gamma_bar(self):
        """
        Return ionization rate (as a function of radius) assuming optical
        depth to cells and of cells is small.
        """
        if not hasattr(self, '_gamma_bar_all'):
            self._gamma_bar_all = \
                np.zeros([self.grid.dims, self.grid.N_absorbers,
                    self.grid.N_absorbers])

            if not self.pf['secondary_ionization']:
                return self._gamma_bar_all

            for i, absorber in enumerate(self.grid.absorbers):
                for j, otherabsorber in enumerate(self.grid.absorbers):
                    self._gamma_bar_all[..., i, j] = self.Gamma_bar[j] \
                        * (self.hnu_bar[j] * self.sigma_tilde[j] \
                        /  self.hnu_bar[i] / self.sigma_bar[j] \
                        - self.grid.ioniz_thresholds[otherabsorber] \
                        / self.grid.ioniz_thresholds[absorber])

        return self._gamma_bar_all

    @property
    def Heat_bar(self):
        """
        Return ionization rate (as a function of radius) assuming optical
        depth to cells and of cells is small.
        """
        if not hasattr(self, '_Heat_bar_all'):
            self._Heat_bar_all = \
                np.zeros([self.grid.dims, self.grid.N_absorbers])
            for i, absorber in enumerate(self.grid.absorbers):
                self._Heat_bar_all[..., i] = self.Gamma_bar[..., i] \
                    * erg_per_ev * (self.hnu_bar[i] * self.sigma_tilde[i] \
                    / self.sigma_bar[i] - self.grid.ioniz_thresholds[absorber])

        return self._Heat_bar_all

    def IonizingPhotonLuminosity(self, t=0, bin=None):
        """
        Return Qdot (photons / s) for this source at energy E.
        """

        if self.pf['source_type'] in [0, 1, 2]:
            return self.Qdot[bin]
        else:
            # Currently only BHs have a time-varying bolometric luminosity
            return self.BolometricLuminosity(t) * self.LE[bin] / self.E[bin] / erg_per_ev

    #def _Intensity(self, E, i, Type, t=0, absorb=True):
    #    """
    #    Return quantity *proportional* to fraction of bolometric luminosity emitted
    #    at photon energy E.  Normalization handled separately.
    #    """
    #
    #    Lnu = self.src._Intensity(E, i, Type, t=t)
    #
    #    # Apply absorbing column
    #    if self.SpectrumPars['logN'][i] > 0 and absorb:
    #        return Lnu * np.exp(-10.**self.SpectrumPars['logN'][i] \
    #            * (sigma_E(E, 0) + y * sigma_E(E, 1)))
    #    else:
    #        return Lnu
    #
    def Spectrum(self, E, t=0.0):
        r"""
        Return fraction of bolometric luminosity emitted at energy E.

        Elsewhere denoted as :math:`I_{\nu}`, normalized such that
        :math:`\int I_{\nu} d\nu = 1`

        Parameters
        ----------
        E: float
            Emission energy in eV
        t: float
            Time in seconds since source turned on.
        i: int
            Index of component to include. If None, includes contribution
            from all components.

        Returns
        -------
        Fraction of bolometric luminosity emitted at E in units of
        eV\ :sup:`-1`\.

        """

        return self._normL * self._Intensity(E, t=t)

    def BolometricLuminosity(self, t=0.0, M=None):
        """
        Returns the bolometric luminosity of a source in units of erg/s.
        For accreting black holes, the bolometric luminosity will increase
        with time, hence the optional 't' and 'M' arguments.
        """

        if self._name == 'bh':
            return self.Luminosity(t, M)
        else:
            return self.Luminosity(t)

    def _FrequencyAveragedBin(self, absorber='h_1', Emin=None, Emax=None,
        energy_weighted=False):
        """
        Bolometric luminosity / number of ionizing photons in spectrum in bandpass
        spanning interval (Emin, Emax). Returns mean photon energy and number of
        ionizing photons in band.
        """

        if Emin is None:
            Emin = max(self.grid.ioniz_thresholds[absorber], self.Emin)
        if Emax is None:
            Emax = self.Emax

        if energy_weighted:
            f = lambda x: x
        else:
            f = lambda x: 1.0

        L = self.Lbol * quad(lambda x: self.Spectrum(x) * f(x), Emin, Emax)[0]
        Q = self.Lbol * quad(lambda x: self.Spectrum(x) * f(x) / x, Emin,
            Emax)[0] / erg_per_ev

        return L / Q / erg_per_ev, Q

    def dump(self, fn, E, clobber=False):
        """
        Write SED out to file.

        Parameters
        ----------
        fn : str
            Filename, suffix determines type. If 'hdf5' or 'h5' will write
            to HDF5 file, otherwise, to ASCII.
        E : np.ndarray
            Array of photon energies at which to sample SED. Units = eV.

        """

        if os.path.exists(fn) and (clobber == False):
            raise OSError('%s exists!')

        if re.search('.hdf5', fn) or re.search('.h5', fn):
            out = 'hdf5'
        else:
            out = 'ascii'

        LE = map(self.Spectrum, E)

        if out == 'hdf5':
            f = h5py.File(fn, 'w')
            f.create_dataset('E', data=E)
            f.create_dataset('LE', data=LE)
            f.close()
        else:
            f = open(fn, 'w')
            print >> f, "# E     LE"
            for i, nrg in enumerate(E):
                print >> f, "%.8e %.8e" % (nrg, LE[i])
            f.close()

        print "Wrote %s." % fn

    def sed_name(self, i=0):
        """
        Return name of output file based on SED properties.
        """

        name = '%s_logM_%.2g_Gamma_%.3g_fsc_%.3g_logE_%.2g-%.2g' % \
            (self.SpectrumPars['type'][i], np.log10(self.src.M0),
             self.src.spec_pars['alpha'][i],
             self.src.spec_pars['fsc'][i], self.logEmin, self.logEmax)

        return name