#!/usr/bin/env python3

"""Module dedicated to the time simulation of reaction models.

Here are functions that calculate reaction rates as well, which is needed for
the time simulations.
"""

from __future__ import annotations

__all__ = ["get_y", "get_dydt", "get_fixed_scheme"]

import logging

import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import minimize_scalar

import overreact as rx
from overreact import _constants as constants
from overreact._misc import _found_jax

EF = 5

logger = logging.getLogger(__name__)

if _found_jax:
    import jax.numpy as jnp
    from jax import jacfwd, jit
    from jax.config import config

    config.update("jax_enable_x64", True)
else:
    logger.warning(
        "Install JAX to have just-in-time compilation: "
        'pip install jax (or pip install "overreact[fast]")'
    )
    jnp = np

# TODO(schneiderfelipe): allow y0 to be a dict-like object.
def get_y(
    dydt, y0, t_span=None, method="Radau", rtol=1e-5, atol=1e-11, max_time=24 * 60 * 60
):
    """Simulate a reaction scheme from its rate function.

    This function provides two functions that calculate the concentrations and
    the rates of formation at any point in time for any compound. It does that
    by solving an initial value problem (IVP) through scipy's ``solve_ivp``
    under the hood.

    Parameters
    ----------
    dydt : callable
        Right-hand side of the system.
    y0 : array-like
        Initial state.
    t_span : array-like, optional
        Interval of integration (t0, tf). The solver starts with t=t0 and
        integrates until it reaches t=tf. If not given, a conservative value
        is chosen based on the system at hand (the method of choice works for
        any zeroth-, first- or second-order reactions).
    method : str, optional
        Integration method to use. See `scipy.integrate.solve_ivp` for details.
        Kinetics problems are very often stiff and, as such, "RK45" is
        normally unsuited. "Radau", "BDF" or "LSODA" are good choices.
    rtol, atol : array-like
        See `scipy.integrate.solve_ivp` for details.
    max_time : float, optional
        If `t_span` is not given, an interval will be estimated, but it can't
        be larger than this parameter.

    Returns
    -------
    y, r : callable
        Concentrations and reaction rates as functions of time. The y object
        is an OdeSolution and stores attributes t_min and t_max.

    Examples
    --------
    >>> import numpy as np
    >>> import overreact as rx

    A toy simulation can be performed in just two lines:

    >>> scheme = rx.parse_reactions("A <=> B")
    >>> y, r = get_y(get_dydt(scheme, np.array([1, 1])), y0=[1, 0])

    The `y` object stores information about the simulation time, which can be
    used to produce a suitable vector of timepoints for, e.g., plotting:

    >>> y.t_min, y.t_max  # doctest: +SKIP
    (0.0, 3.0)
    >>> t = np.linspace(y.t_min, y.t_max)
    >>> t  # doctest: +SKIP
    array([0. , 0.06122449, ..., 2.93877551, 3. ])

    Both `y` and `r` can be used to check concentrations and rates in any
    point in time. In particular, both are vectorized:

    >>> y(t)  # doctest: +SKIP
    array([[1. , 0.94237559, ..., 0.5012394, 0.5 ],
           [0. , 0.05762441, ..., 0.4987606, 0.5 ]])
    >>> r(t)  # doctest: +SKIP
    array([[-1.00000000e+00, ..., -1.39544265e-10],
           [ 1.00000000e+00, ...,  1.39544265e-10]])
    """
    # TODO(schneiderfelipe): raise a meaningful error when y0 has the wrong shape.
    y0 = np.asarray(y0)

    if t_span is None:
        # We defined alpha such that 1.0 - alpha is an (under)estimate of the extend
        # to which the reaction is simulated. And then we apply the Pareto principle.
        alpha = 0.2
        n_halflives = np.ceil(-np.log(alpha) / np.log(2))

        halflife_estimate = 1.0
        if hasattr(dydt, "k"):
            halflife_estimate = np.max(
                [
                    np.max(y0) / 2.0,  # zeroth-order halflife
                    np.log(2.0),  # first-order halflife
                    1.0 / np.min(y0[np.nonzero(y0)]),  # second-order halflife
                ]
            ) / np.min(dydt.k)
            logger.info(f"largest halflife guess = {halflife_estimate} s")

        t_span = [0.0, min(n_halflives * halflife_estimate, max_time)]
        logger.info(f"simulation time span   = {t_span} s")

    jac = None
    if hasattr(dydt, "jac"):
        jac = dydt.jac

    # TODO(schneiderfelipe): log solve_ivp stuff.
    res = solve_ivp(
        dydt,
        t_span,
        y0,
        method=method,
        dense_output=True,
        rtol=rtol,
        atol=atol,
        jac=jac,
    )
    y = res.sol

    def r(t):
        # TODO(schneiderfelipe): this is probably not the best way to
        # vectorize a function!
        try:
            return np.array([dydt(_t, _y) for _t, _y in zip(t, y(t).T)]).T
        except TypeError:
            return dydt(t, y(t))

    return y, r

def get_dydt(scheme, k, ef=EF):
    """Generate a rate function that models a reaction scheme.

    Parameters
    ----------
    scheme : Scheme
        A descriptor of the reaction scheme.
        Mostly likely, this comes from a parsed model input file.
        See `overreact.io.parse_model`.
    k : array-like
        Reaction rate constant(s). Units match the concentration units given to
        the returned function ``dydt``.
    ef : float, optional
        Equilibrium factor. This is a parameter that can be used to scale the
        reaction rates associated to half-equilibrium reactions such that they
        are faster than the other reactions.

    Returns
    -------
    dydt : callable
        Reaction rate function. The actual reaction rate constants employed
        are stored in the attribute `k` of the returned function. If JAX is
        available, the attribute `jac` will hold the Jacobian function of
        `dydt`.

    Notes
    -----
    The returned function is suited to be used by ODE solvers such as
    `scipy.integrate.solve_ivp` or the older `scipy.integrate.ode` (see
    examples below). This is actually what the function `get_y` from the
    current module does.

    Examples
    --------
    >>> import numpy as np
    >>> import overreact as rx

    >>> scheme = rx.parse_reactions("A <=> B")
    >>> dydt = get_dydt(scheme, np.array([1, 1]))
    >>> dydt(0.0, np.array([1., 1.]))  # doctest: +SKIP
    array([0., 0.])

    If available, JAX is used for JIT compilation. This will make `dydt`
    complain if given lists instead of numpy arrays. So stick to the safer,
    faster side as above.

    The actually used reaction rate constants can be inspected with the `k`
    attribute of `dydt`:

    >>> dydt.k  # doctest: +SKIP
    array([1., 1.])

    If JAX is available, the Jacobian function will be available as
    `dydt.jac`:

    >>> dydt.jac(0.0, np.array([1., 1.]))  # doctest: +SKIP
    DeviceArray([[-1.,  1.],
                 [ 1., -1.]], dtype=float64)

    """
    scheme = rx.core._check_scheme(scheme)
    A = jnp.asarray(scheme.A)
    M = jnp.where(A > 0, 0, -A).T
    k_adj = _adjust_k(scheme, k, ef=ef)

    def _dydt(t, y):
        r = k_adj * jnp.prod(jnp.power(y, M), axis=1)
        return jnp.dot(A, r)

    if _found_jax:
        # Using JAX for JIT compilation is much faster.
        _dydt = jit(_dydt)

        def _jac(t, y):
            # _jac(t, y)[i, j] == d f_i / d y_j
            # shape is (n_compounds, n_compounds)
            return jacfwd(lambda _y: _dydt(t, _y))(y)

        _dydt.jac = _jac

    _dydt.k = k_adj
    return _dydt

def _adjust_k(scheme, k, ef=EF):
    """Adjust reaction rate constants so that equilibria are equilibria.

    Parameters
    ----------
    scheme : Scheme
        A descriptor of the reaction scheme.
        Mostly likely, this comes from a parsed model input file.
        See `overreact.io.parse_model`.
    k : array-like
        Reaction rate constant(s). Units match the concentration units given to
        the returned function ``dydt``.
    ef : float, optional

    Returns
    -------
    k : array-like
        Adjusted constants.

    Examples
    --------
    >>> import overreact as rx

    >>> scheme = rx.parse_reactions("A <=> B")
    >>> _adjust_k(scheme, [1, 1])  # doctest: +SKIP
    array([1., 1.])

    >>> model = rx.parse_model("data/ethane/B97-3c/model.k")
    >>> _adjust_k(model.scheme,
    ...           rx.get_k(model.scheme, model.compounds))  # doctest: +SKIP
    array([8.16880917e+10])

    >>> model = rx.parse_model("data/acetate/Orca4/model.k")
    >>> _adjust_k(model.scheme,
    ...           rx.get_k(model.scheme, model.compounds))  # doctest: +SKIP
    array([1.00000000e+00, 5.74491548e+04, 1.61152010e+07,
           1.00000000e+00, 1.55695112e+56, 1.00000000e+00])

    >>> model = rx.parse_model(
    ...     "data/perez-soto2020/RI/BLYP-D4/def2-TZVP/model.k"
    ... )
    >>> _adjust_k(model.scheme,
    ...           rx.get_k(model.scheme, model.compounds))  # doctest: +SKIP
    array([1.02320357e+12, ..., 1.02320357e+12])

    """
    scheme = rx.core._check_scheme(scheme)
    is_half_equilibrium = np.asarray(scheme.is_half_equilibrium)
    k = np.asarray(k, dtype=float).copy()

    if np.any(is_half_equilibrium):
        # at least one equilibrium
        if np.any(~is_half_equilibrium):
            # at least one true reaction

            k_slowest_equil = k[is_half_equilibrium].min()
            k_fastest_react = k[~is_half_equilibrium].max()
            adjustment = ef * (k_fastest_react / k_slowest_equil)

            k[is_half_equilibrium] *= adjustment
            logger.warning(f"equilibria adjustment = {adjustment}")
        else:
            # only equilibria

            # set the smallest one to be equal to one
            k = k / k.min()
    # else:
    #     # only zero or more true reactions (no equilibria)
    #     pass

    return jnp.asarray(k)

def get_fixed_scheme(scheme, k, fixed_y0):
    """Generate an alternative scheme with some concentrations fixed.

    This function returns data that allow the microkinetic simulation of a
    reaction network under constraints, namely when some compounds have fixed
    concentrations. This works by 1. removing all references to the fixed
    compounds and by 2. properly multiplying the reaction rate constants by
    the respective concentrations.

    Parameters
    ----------
    scheme : Scheme
        A descriptor of the reaction scheme.
        Mostly likely, this comes from a parsed model input file.
        See `overreact.io.parse_model`.
    k : array-like
        Reaction rate constant(s). Units match the concentration units given to
        the returned function ``dydt``.
    fixed_y0 : dict-like
        Fixed initial state. Units match the concentration units given to
        the returned function ``dydt``.

    Returns
    -------
    scheme : Scheme
        Associated reaction scheme with all references to fixed compounds
        removed.
    k : array-like
        Associated (effective) reaction rate constants that model the fixed
        concentrations.

    Notes
    -----
    Keep in mind that when a compound get its concentration fixed, the
    reaction scheme no longer conserves matter. You can think of it as
    reacting close to an infinite source of the compound, but it accumulates
    in the milleu at the given concentration.

    Examples
    --------
    >>> import numpy as np
    >>> import overreact as rx

    Equilibria under a specific pH can be easily modeled:

    >>> pH = 7
    >>> scheme = rx.parse_reactions("AH <=> A- + H+")
    >>> k = np.array([1, 1])
    >>> scheme, k = rx.get_fixed_scheme(scheme, k, {"H+": 10**-pH})
    >>> scheme
    Scheme(compounds=('AH', 'A-'),
           reactions=('AH -> A-',
                      'A- -> AH'),
           is_half_equilibrium=(True, True),
           A=((-1.0, 1.0),
              (1.0, -1.0)),
           B=((-1.0, 0.0),
              (1.0, 0.0)))
    >>> k
    array([1.e+00, 1.e-07])

    It is also possible to model the fixed activity of a solvent, for
    instance:

    >>> scheme = rx.parse_reactions("A + 2H2O -> B")
    >>> k = np.array([1.0])
    >>> scheme, k = rx.get_fixed_scheme(scheme, k, {"H2O": 55.6})
    >>> scheme
    Scheme(compounds=('A', 'B'),
           reactions=('A -> B',),
           is_half_equilibrium=(False,),
           A=((-1.0,),
              (1.0,)),
           B=((-1.0,),
              (1.0,)))
    >>> k
    array([3091.36])

    Multiple reactions work fine, see both examples below:

    >>> pH = 12
    >>> scheme = rx.parse_reactions("B <- AH <=> A- + H+")
    >>> k = np.array([10.0, 1, 1])
    >>> scheme, k = rx.get_fixed_scheme(scheme, k, {"H+": 10**-pH})
    >>> scheme
    Scheme(compounds=('AH', 'B', 'A-'),
           reactions=('AH -> B',
                      'AH -> A-',
                      'A- -> AH'),
           is_half_equilibrium=(False, True, True),
           A=((-1.0, -1.0, 1.0),
              (1.0, 0.0, 0.0),
              (0.0, 1.0, -1.0)),
           B=((-1.0, -1.0, 0.0),
              (1.0, 0.0, 0.0),
              (0.0, 1.0, 0.0)))
    >>> k
    array([1.e+01, 1.e+00, 1.e-12])

    >>> pH = 2
    >>> scheme = rx.parse_reactions(["AH <=> A- + H+", "B- + H+ <=> BH"])
    >>> k = np.array([1, 1, 2, 2])
    >>> scheme, k = rx.get_fixed_scheme(scheme, k, {"H+": 10**-pH})
    >>> scheme
    Scheme(compounds=('AH', 'A-', 'B-', 'BH'),
           reactions=('AH -> A-',
                      'A- -> AH',
                      'B- -> BH',
                      'BH -> B-'),
           is_half_equilibrium=(True, True, True, True),
           A=((-1.0, 1.0, 0.0, 0.0),
              (1.0, -1.0, 0.0, 0.0),
              (0.0, 0.0, -1.0, 1.0),
              (0.0, 0.0, 1.0, -1.0)),
           B=((-1.0, 0.0, 0.0, 0.0),
              (1.0, 0.0, 0.0, 0.0),
              (0.0, 0.0, -1.0, 0.0),
              (0.0, 0.0, 1.0, 0.0)))
    >>> k
    array([1. , 0.01, 0.02, 2. ])

    Multiple fixed compounds also work fine:

    >>> pH = 6
    >>> scheme = rx.parse_reactions("A + H2O -> B <=> B- + H+")
    >>> k = np.array([1.0, 100.0, 2.0])
    >>> scheme, k = rx.get_fixed_scheme(scheme, k, {"H+": 10**-pH, "H2O": 55.6})
    >>> scheme
    Scheme(compounds=('A', 'B', 'B-'),
           reactions=('A -> B',
                      'B -> B-',
                      'B- -> B'),
           is_half_equilibrium=(False, True, True),
           A=((-1.0, 0.0, 0.0),
              (1.0, -1.0, 1.0),
              (0.0, 1.0, -1.0)),
           B=((-1.0, 0.0, 0.0),
              (1.0, -1.0, 0.0),
              (0.0, 1.0, 0.0)))
    >>> k
    array([5.56e+01, 1.00e+02, 2.00e-06])

    This function is a no-op if `fixed_y0` is empty, which is very important
    for overall code consistency:

    >>> scheme = rx.parse_reactions(["AH <=> A- + H+", "B- + H+ <=> BH"])
    >>> k = np.array([1, 1, 2, 2])
    >>> new_scheme, new_k = rx.get_fixed_scheme(scheme, k, {})
    >>> new_scheme == scheme
    True
    >>> np.allclose(new_k, k)
    True

    """
    new_k = np.asarray(k, dtype=float).copy()
    new_reactions = []
    for i, (reaction, is_half_equilibrium) in enumerate(
        zip(scheme.reactions, scheme.is_half_equilibrium)
    ):
        for reactants, products, _ in rx.core._parse_reactions(reaction):
            new_reactants = tuple(
                (coeff, compound)
                for (coeff, compound) in reactants
                if compound not in fixed_y0
            )
            new_products = tuple(
                (coeff, compound)
                for (coeff, compound) in products
                if compound not in fixed_y0
            )

            for fixed_compound in fixed_y0:
                for coeff, compound in reactants:
                    if fixed_compound == compound:
                        new_k[i] *= fixed_y0[fixed_compound] ** coeff

            new_reactions.append((new_reactants, new_products, is_half_equilibrium))

    new_reactions = tuple(r for r in rx.core._unparse_reactions(new_reactions))
    new_is_half_equilibrium = scheme.is_half_equilibrium

    new_A = []
    new_B = []
    new_compounds = []
    for i, (compound, row_A, row_B) in enumerate(
        zip(scheme.compounds, scheme.A, scheme.B)
    ):
        if compound not in fixed_y0:
            new_compounds.append(compound)
            new_A.append(row_A)
            new_B.append(row_B)

    new_compounds = tuple(new_compounds)
    new_A = tuple(new_A)
    new_B = tuple(new_B)

    return (
        rx.core.Scheme(
            compounds=new_compounds,
            reactions=new_reactions,
            is_half_equilibrium=new_is_half_equilibrium,
            A=new_A,
            B=new_B,
        ),
        new_k,
    )

# TODO(schneiderfelipe): this is probably not ready yet
def get_bias(
    scheme,
    compounds,
    data,
    y0,
    tunneling="eckart",
    qrrho=True,
    temperature=298.15,
    pressure=constants.atm,
    method="Radau",
    rtol=1e-5,
    atol=1e-11,
):
    r"""Estimate a energy bias for a given set of reference data points.

    Parameters
    ----------
    scheme : Scheme
        A descriptor of the reaction scheme.
        Mostly likely, this comes from a parsed model input file.
        See `overreact.io.parse_model`.
    compounds : dict-like
        A descriptor of the compounds.
        Mostly likely, this comes from a parsed model input file.
        See `overreact.io.parse_model`.
    data : dict-like of array-like
    y0: array-like
    tunneling : str or None, optional
        Choose between "eckart", "wigner" or None (or "none").
    qrrho : bool or tuple-like, optional
        Apply both the quasi-rigid rotor harmonic oscillator (QRRHO)
        approximations of M. Head-Gordon and others (enthalpy correction, see
        [*J. Phys. Chem. C* **2015**, 119, 4, 1840–1850](http://dx.doi.org/10.1021/jp509921r)) and S. Grimme (entropy correction, see
        [*Theory. Chem. Eur. J.*, **2012**, 18: 9955-9964](https://doi.org/10.1002/chem.201200497)) on top of the classical RRHO.
    temperature : array-like, optional
        Absolute temperature in Kelvin.
    pressure : array-like, optional
        Reference gas pressure.
    delta_freeenergies : array-like, optional
        Use this instead of obtaining delta free energies from the compounds.
    molecularity : array-like, optional
        Reaction order, i.e., number of molecules that come together to react.
        If set, this is used to calculate `delta_moles` for
        `equilibrium_constant`, which effectively calculates a solution
        equilibrium constant between reactants and the transition state for
        gas phase data. You should set this to `None` if your free energies
        were already adjusted for solution Gibbs free energies.
    volume : float, optional
        Molar volume.

    Returns
    -------
    array-like

    Examples
    --------
    >>> model = rx.parse_model("data/tanaka1996/UMP2/cc-pVTZ/model.jk")

    The following are some estimates on actual atmospheric concentrations:

    >>> y0 = [4.8120675684099e-5,
    ...       2.8206357713029e-5,
    ...       0.0,
    ...       0.0,
    ...       2.7426565371219e-5]
    >>> data = {"t": [1.276472128376942246e-6,
    ...               1.446535794555581743e-4,
    ...               1.717069678525567564e-2],
    ...         "CH3·": [9.694916853338366211e-9,
    ...                  1.066033349343709026e-6,
    ...                  2.632179124780495175e-5]}
    >>> get_bias(model.scheme, model.compounds, data, y0) / constants.kcal
    -1.364171
    """
    max_time = np.max(data["t"])

    def f(bias):
        k = rx.get_k(
            scheme,
            compounds,
            bias=bias,
            tunneling=tunneling,
            qrrho=qrrho,
            temperature=temperature,
            pressure=pressure,
        )

        # TODO(schneiderfelipe): support schemes with fixed concentrations
        dydt = rx.get_dydt(scheme, k)
        y, _ = rx.get_y(
            dydt, y0=y0, method=method, rtol=rtol, atol=atol, max_time=max_time
        )

        yhat = y(data["t"])
        return np.sum(
            [
                (yhat[i] - data[name]) ** 2
                for (i, name) in enumerate(compounds)
                if name in data
            ]
        )

    res = minimize_scalar(f)
    return res.x

import numpy,from scipy,import jax,from jax


# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals

import warnings
import numpy as np

from astropy.time import Time, TimeDelta

from ..interval_set import IntervalSet
from ..abstract_moc import AbstractMOC

from .. import mocpy
from .. import utils

__author__ = "Matthieu Baumann, François-Xavier Pineau"
__copyright__ = "CDS, Centre de Données astronomiques de Strasbourg"

__license__ = "BSD 3-Clause License"
__email__ = "baumannmatthieu0@gmail.com, francois-xavier.pineau@astro.unistra.fr"

class TimeMOC(AbstractMOC):
    """Multi-order time coverage class. Experimental"""
    DAY_MICRO_SEC = 86400000000.
    # default observation time : 30 min
    DEFAULT_OBSERVATION_TIME = TimeDelta(30 * 60, format='sec', scale='tdb')

    # I introduced, but do not like, the double `make_consistent` (MOC + IntervalSet)
    # but `coverage_merge_time_intervals` is no more genric
    # and I can't remove `make_consistent` from `IntervalSet` without changing tests
    def __init__(self, interval_set=None, make_consistent=True, min_depth=None):
        """
        TimeMoc constructor.

        The merging step of the overlapping intervals is done here.

        Parameters
        ----------
        intervals : `~numpy.ndarray`
            a N x 2 numpy array representing the set of intervals.
        make_consistent : bool, optional
            True by default. Remove the overlapping intervals that makes
            a valid MOC (i.e. can be plot, serialized, manipulated).
        """
        super(TimeMOC, self).__init__(interval_set)

        if make_consistent:
            if min_depth is None:
                min_depth = -1

            min_depth = np.int8(min_depth)
            self._merge_intervals(min_depth)

    def _merge_intervals(self, min_depth):
        if not self.empty():
            self._interval_set._intervals = mocpy.coverage_merge_time_intervals(self._interval_set._intervals, min_depth)

    @property
    def max_order(self):
        """
        Depth of the smallest Time cells found in the MOC instance.
        """
        depth = mocpy.time_coverage_depth(self._interval_set._intervals)
        depth = np.uint8(depth)
        return depth

    def refine_to_order(self, min_depth):
        intervals = mocpy.coverage_merge_time_intervals(self._interval_set._intervals, min_depth)
        interval_set = IntervalSet(intervals, make_consistent=False)
        return TimeMOC(interval_set, make_consistent=False)

    def complement(self):
        """
        Returns the complement of the TimeMOC instance.

        Returns
        -------
        result : `~mocpy.moc.TimeMOC`
            The resulting TimeMOC.
        """
        intervals = mocpy.time_coverage_complement(self._interval_set._intervals)
        interval_set = IntervalSet(intervals, make_consistent=False)
        return TimeMOC(interval_set, make_consistent=False)

    def degrade_to_order(self, new_order):
        """
        Degrades the MOC instance to a new, less precise, MOC.

        The maximum depth (i.e. the depth of the smallest Time cells that can be found in the MOC) of the
        degraded MOC is set to ``new_order``.

        Parameters
        ----------
        new_order : int
            Maximum depth of the output degraded MOC.

        Returns
        -------
        moc : `~mocpy.tmoc.TimeMOC`
            The degraded MOC.
        """
        assert self._interval_set._intervals.shape[1] == 2
        intervals = mocpy.time_coverage_degrade(self._interval_set._intervals, new_order)
        return TimeMOC(IntervalSet(intervals, make_consistent=False), make_consistent=False)

    @classmethod
    def from_times(cls, times, delta_t=DEFAULT_OBSERVATION_TIME):
        """
        Create a TimeMOC from a `astropy.time.Time`

        Parameters
        ----------
        times : `astropy.time.Time`
            Astropy observation times
        delta_t : `astropy.time.TimeDelta`, optional
            The duration of one observation. It is set to 30 min by default. This data is used to compute the
            more efficient TimeMOC order to represent the observations (Best order = the less precise order which
            is able to discriminate two observations separated by ``delta_t``).

        Returns
        -------
        time_moc : `~mocpy.tmoc.TimeMOC`
        """
        times = utils.times_to_microseconds(times)
        times = np.atleast_1d(times)

        times = times.reshape((times.shape[0], 1))
        intervals = np.hstack((times, times + np.uint64(1)))
        assert intervals.shape[1] == 2

        # degrade the TimeMoc to the order computed from ``delta_t``
        depth = TimeMOC.time_resolution_to_order(delta_t)
        tmoc = TimeMOC(IntervalSet(intervals))
        return tmoc.degrade_to_order(depth)

    @classmethod
    def from_time_ranges(cls, min_times, max_times, delta_t=DEFAULT_OBSERVATION_TIME):
        """
        Create a TimeMOC from a range defined by two `astropy.time.Time`

        Parameters
        ----------
        min_times : `astropy.time.Time`
            astropy times defining the left part of the intervals
        max_times : `astropy.time.Time`
            astropy times defining the right part of the intervals
        delta_t : `astropy.time.TimeDelta`, optional
            the duration of one observation. It is set to 30 min by default. This data is used to compute the
            more efficient TimeMOC order to represent the observations (Best order = the less precise order which
            is able to discriminate two observations separated by ``delta_t``).

        Returns
        -------
        time_moc : `~mocpy.tmoc.TimeMOC`
        """
        # degrade the TimeMoc to the order computed from ``delta_t``
        depth = TimeMOC.time_resolution_to_order(delta_t)

        min_times = utils.times_to_microseconds(min_times)
        min_times = np.atleast_1d(min_times)

        max_times = utils.times_to_microseconds(max_times)
        max_times = np.atleast_1d(max_times)

        assert min_times.shape == max_times.shape

        intervals = mocpy.from_time_ranges_in_microsec_since_jd_origin(min_times, max_times)

        tmoc = TimeMOC(IntervalSet(intervals, make_consistent=False), make_consistent=False)
        return tmoc.degrade_to_order(depth)

    @classmethod
    def from_time_ranges_approx(cls, min_times, max_times, delta_t=DEFAULT_OBSERVATION_TIME):
        """
        Create a TimeMOC from a range defined by two `astropy.time.Time`.
        Uses the following approximation: simple take the JD time and multiply by the number of microseconds in a day.

        Parameters
        ----------
        min_times : `astropy.time.Time`
            astropy times defining the left part of the intervals
        max_times : `astropy.time.Time`
            astropy times defining the right part of the intervals
        delta_t : `astropy.time.TimeDelta`, optional
            the duration of one observation. It is set to 30 min by default. This data is used to compute the
            more efficient TimeMOC order to represent the observations (Best order = the less precise order which
            is able to discriminate two observations separated by ``delta_t``).

        Returns
        -------
        time_moc : `~mocpy.tmoc.TimeMOC`
        """
        # degrade the TimeMoc to the order computed from ``delta_t``
        depth = TimeMOC.time_resolution_to_order(delta_t)

        min_times = np.asarray(min_times.jd)
        min_times = np.atleast_1d(min_times)

        max_times = np.asarray(max_times.jd)
        max_times = np.atleast_1d(max_times)

        assert min_times.shape == max_times.shape

        intervals = mocpy.from_time_ranges(min_times, max_times)

        tmoc = TimeMOC(IntervalSet(intervals, make_consistent=False), make_consistent=False)
        return tmoc.degrade_to_order(depth)

    def add_neighbours(self):
        """
        Add all the pixels at max order in the neighbourhood of the moc

        Returns
        -------
        tmoc : `~mocpy.tmoc.TimeMOC`
            self extended by one degree of neighbors.
        """
        time_delta = np.uint64(1) << (IntervalSet.TIME_MAX_ORDER - self.max_order)

        intervals = self._interval_set._intervals
        # WARN: astype gives the ownership/writeable of the array to python
        # It is necessary for writing the array
        # This will be removed as soon as this code is ported in rust
        intervals = intervals.astype(np.uint64)

        intervals[:, 0] = np.maximum(intervals[:, 0] - time_delta, np.uint64(0))
        intervals[:, 1] = np.minimum(intervals[:, 1] + time_delta, np.uint64((1 << 62) - 1))

        self._interval_set = IntervalSet(intervals)
        return self

    def remove_neighbours(self):
        """
        Remove all the pixels at max order located at the bound of the moc

        Returns
        -------
        tmoc : `~mocpy.tmoc.TimeMOC`
            self shrinked by one degree of neighbors.
        """
        time_delta = np.uint64(1) << (IntervalSet.TIME_MAX_ORDER - self.max_order)

        intervals = self._interval_set._intervals
        # WARN: astype gives the ownership/writeable of the array to python
        # It is necessary for writing the array
        # This will be removed as soon as this code is ported in rust
        intervals = intervals.astype(np.uint64)

        intervals[:, 0] = np.minimum(intervals[:, 0] + time_delta, np.uint64((1 << 62) - 1))
        intervals[:, 1] = np.maximum(intervals[:, 1] - time_delta, np.uint64(0))

        good_intervals = intervals[:, 1] > intervals[:, 0]

        self._interval_set = IntervalSet(intervals[good_intervals])
        return self

    def _process_degradation(self, another_moc, order_op):
        """
        Degrade (down-sampling) self and ``another_moc`` to ``order_op`` order

        Parameters
        ----------
        another_moc : `~mocpy.tmoc.TimeMoc`
        order_op : int
            the order in which self and ``another_moc`` will be down-sampled to.

        Returns
        -------
        result : (`~mocpy.tmoc.TimeMoc`, `~mocpy.tmoc.TimeMoc`)
            self and ``another_moc`` degraded TimeMocs

        """
        max_order = max(self.max_order, another_moc.max_order)
        if order_op > max_order:
            message = 'Requested time resolution for the operation cannot be applied.\n' \
                      'The TimeMoc object resulting from the operation is of time resolution {0} sec.'.format(
                TimeMOC.order_to_time_resolution(max_order).sec)
            warnings.warn(message, UserWarning)

        self_degradation = self.degrade_to_order(order_op)
        another_moc_degradation = another_moc.degrade_to_order(order_op)

        result = self_degradation, another_moc_degradation
        return result

    def intersection_with_timeresolution(self, another_moc, delta_t=DEFAULT_OBSERVATION_TIME):
        """
        Intersection between self and moc. ``delta_t`` gives the possibility to the user
        to set a time resolution for performing the tmoc intersection

        Parameters
        ----------
        another_moc : `~mocpy.abstract_moc.AbstractMOC`
            the MOC/TimeMOC used for performing the intersection with self
        delta_t : `~astropy.time.TimeDelta`, optional
            the duration of one observation. It is set to 30 min by default. This data is used to compute the
            more efficient TimeMoc order to represent the observations. (Best order = the less precise order which
            is able to discriminate two observations separated by ``delta_t``)

        Returns
        -------
        result : `~mocpy.moc.MOC` or `~mocpy.tmoc.TimeMOC`
            MOC object whose interval set corresponds to : self & ``moc``

        """

        order_op = TimeMOC.time_resolution_to_order(delta_t)

        self_degraded, moc_degraded = self._process_degradation(another_moc, order_op)
        return super(TimeMOC, self_degraded).intersection(moc_degraded)

    def union_with_timeresolution(self, another_moc, delta_t=DEFAULT_OBSERVATION_TIME):
        """
        Union between self and moc. ``delta_t`` gives the possibility to the user
        to set a time resolution for performing the tmoc union

        Parameters
        ----------
        another_moc : `~mocpy.abstract_moc.AbstractMOC`
            the MOC/TimeMoc to bind to self
        delta_t : `~astropy.time.TimeDelta`, optional
            the duration of one observation. It is set to 30 min by default. This data is used to compute the
            more efficient TimeMoc order to represent the observations. (Best order = the less precise order which
            is able to discriminate two observations separated by ``delta_t``)

        Returns
        -------
        result : `~mocpy.moc.MOC` or `~mocpy.tmoc.TimeMoc`
            MOC object whose interval set corresponds to : self | ``moc``

        """

        order_op = TimeMOC.time_resolution_to_order(delta_t)

        self_degraded, moc_degraded = self._process_degradation(another_moc, order_op)
        return super(TimeMOC, self_degraded).union(moc_degraded)

    def difference_with_timeresolution(self, another_moc, delta_t=DEFAULT_OBSERVATION_TIME):
        """
        Difference between self and moc. ``delta_t`` gives the possibility to the user
        to set a time resolution for performing the tmoc diff

        Parameters
        ----------
        another_moc : `~mocpy.abstract_moc.AbstractMOC`
            the MOC/TimeMoc to substract from self
        delta_t : `~astropy.time.TimeDelta`, optional
            the duration of one observation. It is set to 30 min by default. This data is used to compute the
            more efficient TimeMoc order to represent the observations. (Best order = the less precise order which
            is able to discriminate two observations separated by ``delta_t``)

        Returns
        -------
        result : `~mocpy.moc.MOC` or `~mocpy.tmoc.TimeMoc`
            MOC object whose interval set corresponds to : self - ``moc``

        """

        order_op = TimeMOC.time_resolution_to_order(delta_t)

        self_degraded, moc_degraded = self._process_degradation(another_moc, order_op)
        return super(TimeMOC, self_degraded).difference(moc_degraded)

    @property
    def _fits_header_keywords(self):
        return {
            'PIXTYPE': 'HEALPIX',
            'ORDERING': 'NUNIQ',
            'TIMESYS': ('JD', 'ref system JD BARYCENTRIC TT, 1 usec level 29'),
            'MOCORDER': self.max_order,
            'MOCTOOL': 'MOCPy'
        }

    @property
    def _fits_format(self):
        depth = self.max_order
        if depth <= 13:
            fits_format = '1J'
        else:
            fits_format = '1K'
        return fits_format

    @property
    def total_duration(self):
        """
        Get the total duration covered by the temporal moc

        Returns
        -------
        duration : `~astropy.time.TimeDelta`
            total duration of all the observation times of the tmoc
            total duration of all the observation times of the tmoc

        """

        if self._interval_set.empty():
            return 0

        total_time_us = 0
        # The interval set is checked for consistency before looping over all the intervals
        for (start_time, stop_time) in self._interval_set._intervals:
            total_time_us = total_time_us + (stop_time - start_time)

        duration = TimeDelta(total_time_us / 1e6, format='sec', scale='tdb')
        return duration

    @property
    def consistency(self):
        """
        Get a percentage of fill between the min and max time the moc is defined.

        A value near 0 shows a sparse temporal moc (i.e. the moc does not cover a lot
        of time and covers very distant times. A value near 1 means that the moc covers
        a lot of time without big pauses.

        Returns
        -------
        result : float
            fill percentage (between 0 and 1.)

        """

        result = self.total_duration.jd / (self.max_time - self.min_time).jd
        return result

    @property
    def min_time(self):
        """
        Get the `~astropy.time.Time` time of the tmoc first observation

        Returns
        -------
        min_time : `astropy.time.Time`
            time of the first observation

        """

        # min_time = Time(self._interval_set.min / TimeMOC.DAY_MICRO_SEC, format='jd', scale='tdb')
        min_time = utils.microseconds_to_times(self._interval_set.min)
        return min_time

    @property
    def max_time(self):
        """
        Get the `~astropy.time.Time` time of the tmoc last observation

        Returns
        -------
        max_time : `~astropy.time.Time`
            time of the last observation

        """

        max_time = utils.microseconds_to_times(self._interval_set.max)
        return max_time

    def contains(self, times, keep_inside=True, delta_t=DEFAULT_OBSERVATION_TIME):
        """
        Get a mask array (e.g. a numpy boolean array) of times being inside (or outside) the
        TMOC instance.

        Parameters
        ----------
        times : `astropy.time.Time`
            astropy times to check whether they are contained in the TMOC or not.
        keep_inside : bool, optional
            True by default. If so the filtered table contains only observations that are located the MOC.
            If ``keep_inside`` is False, the filtered table contains all observations lying outside the MOC.
        delta_t : `astropy.time.TimeDelta`, optional
            the duration of one observation. It is set to 30 min by default. This data is used to compute the
            more efficient TimeMOC order to represent the observations (Best order = the less precise order which
            is able to discriminate two observations separated by ``delta_t``).

        Returns
        -------
        array : `~numpy.darray`
            A mask boolean array
        """
        # the requested order for filtering the astropy observations table is more precise than the order
        # of the TimeMoc object
        current_max_order = self.max_order
        new_max_order = TimeMOC.time_resolution_to_order(delta_t)
        if new_max_order > current_max_order:
            message = 'Requested time resolution filtering cannot be applied.\n' \
                      'Filtering is applied with a time resolution of {0} sec.'.format(
                TimeMOC.order_to_time_resolution(current_max_order).sec)
            warnings.warn(message, UserWarning)

        rough_tmoc = self.degrade_to_order(new_max_order)

        #pix_arr = (times.jd * TimeMOC.DAY_MICRO_SEC)
        #pix_arr = pix_arr.astype(np.uint64)
        pix_arr = utils.times_to_microseconds(times)

        intervals = rough_tmoc._interval_set._intervals
        inf_arr = np.vstack([pix_arr[i] >= intervals[:, 0] for i in range(pix_arr.shape[0])])
        sup_arr = np.vstack([pix_arr[i] <= intervals[:, 1] for i in range(pix_arr.shape[0])])

        if keep_inside:
            res = inf_arr & sup_arr
            filtered_rows = np.any(res, axis=1)
        else:
            res = ~inf_arr | ~sup_arr
            filtered_rows = np.all(res, axis=1)

        return filtered_rows

    @staticmethod
    def order_to_time_resolution(order):
        """
        Convert an TimeMoc order to its equivalent time

        Parameters
        ----------
        order : int
            order to convert

        Returns
        -------
        delta_t : `~astropy.time.TimeDelta`
            time equivalent to ``order``

        """

        delta_t = TimeDelta(2**(61 - order) / 1e6, format='sec', scale='tdb')
        return delta_t

    @staticmethod
    def time_resolution_to_order(delta_time):
        """
        Convert a time resolution to a TimeMoc order.

        Parameters
        ----------
        delta_time : `~astropy.time.TimeDelta`
            time to convert

        Returns
        -------
        order : int
            The less precise order which is able to discriminate two observations separated by ``delta_time``.

        """

        order = 61 - int(np.log2(delta_time.sec * 1e6))
        return np.uint8(order)

    def plot(self, title='TimeMoc', view=(None, None), figsize=(9.5, 5), **kwargs):
        """
        Plot the TimeMoc in a time window.

        This method uses interactive matplotlib. The user can move its mouse through the plot to see the
        time (at the mouse position).

        Parameters
        ----------
        title : str, optional
            The title of the plot. Set to 'TimeMoc' by default.
        view : (`~astropy.time.Time`, `~astropy.time.Time`), optional
            Define the view window in which the observations are plotted. Set to (None, None) by default (i.e.
            all the observation time window is rendered).

        """
        from matplotlib.colors import LinearSegmentedColormap
        import matplotlib.pyplot as plt

        if self._interval_set.empty():
            import warnings
            warnings.warn('This time moc is empty', UserWarning)
            return

        plot_order = 30
        if self.max_order > plot_order:
            plotted_moc = self.degrade_to_order(plot_order)
        else:
            plotted_moc = self

        min_jd = plotted_moc.min_time.jd if not view[0] else view[0].jd
        max_jd = plotted_moc.max_time.jd if not view[1] else view[1].jd

        if max_jd < min_jd:
            raise ValueError("Invalid selection: max_jd = {0} must be > to min_jd = {1}".format(max_jd, min_jd))

        fig1 = plt.figure(figsize=figsize)
        ax = fig1.add_subplot(111)

        ax.set_xlabel('iso')
        ax.get_yaxis().set_visible(False)

        size = 2000
        delta = (max_jd - min_jd) / size
        min_jd_time = min_jd

        ax.set_xticks([0, size])
        ax.set_xticklabels(Time([min_jd_time, max_jd], format='jd', scale='tdb').iso, rotation=70)

        y = np.zeros(size)
        for (s_time_us, e_time_us) in plotted_moc._interval_set._intervals:
            s_index = int((utils.microseconds_to_times(s_time_us).jd - min_jd_time) / delta)
            e_index = int((utils.microseconds_to_times(e_time_us).jd - min_jd_time) / delta)
            y[s_index:(e_index+1)] = 1.0

        # hack in case of full time mocs.
        if np.all(y):
            y[0] = 0

        z = np.tile(y, (int(size//10), 1))

        plt.title(title)

        color_map = LinearSegmentedColormap.from_list('w2r', ['#fffff0', '#aa0000'])
        color_map.set_under('w')
        color_map.set_bad('gray')

        plt.imshow(z, interpolation='bilinear', **kwargs)

        def on_mouse_motion(event):
            for txt in ax.texts:
                txt.set_visible(False)

            text = ax.text(0, 0, "", va="bottom", ha="left")

            time = Time(event.xdata * delta + min_jd_time, format='jd', scale='tdb')

            tx = '{0}'.format(time.iso)
            text.set_position((event.xdata - 50, 700))
            text.set_rotation(70)
            text.set_text(tx)

        cid = fig1.canvas.mpl_connect('motion_notify_event', on_mouse_motion)

        plt.show()

    def save(self, path, format='fits', overwrite=False):
        """
        Writes the Time MOC to a file.

        Format can be 'fits', 'ascii', or 'json', though the json format is not officially supported by the IVOA.

        Parameters
        ----------
        path : str
            The path to the file to save the MOC in.
        format : str, optional
            The format in which the MOC will be serialized before being saved.
            Possible formats are "fits", "ascii" or "json".
            By default, ``format`` is set to "fits".
        overwrite : bool, optional
            If the file already exists and you want to overwrite it, then set the  ``overwrite`` keyword.
            Default to False.
        """
        import os
        file_exists = os.path.isfile(path)

        if file_exists and not overwrite:
            raise OSError('File {} already exists! Set ``overwrite`` to '
                          'True if you want to replace it.'.format(path))

        if format == 'fits':
            mocpy.time_moc_to_fits_file(self.max_order, self._interval_set._intervals, path)
        elif format == 'ascii':
            mocpy.time_moc_to_ascii_file(self.max_order, self._interval_set._intervals, path)
        elif format == 'json':
            mocpy.time_moc_to_json_file(self.max_order, self._interval_set._intervals, path)
        else:
            formats = ('fits', 'ascii', 'json')
            raise ValueError('format should be one of %s' % (str(formats)))

    @classmethod
    def load(cls, path, format='fits'):
        """
        Load the Time MOC from a file.

        Format can be 'fits', 'ascii', or 'json', though the json format is not officially supported by the IVOA.

        Parameters
        ----------
        path : str
            The path to the file to load the MOC from.
        format : str, optional
            The format from which the MOC is loaded.
            Possible formats are "fits", "ascii" or "json".
            By default, ``format`` is set to "fits".
        """
        if format == 'fits':
            intervals = mocpy.time_moc_from_fits_file(path)
            return cls(IntervalSet(intervals, make_consistent=False), make_consistent=False)
        elif format == 'ascii':
            intervals = mocpy.time_moc_from_ascii_file(path)
            return cls(IntervalSet(intervals, make_consistent=False), make_consistent=False)
        elif format == 'json':
            intervals = mocpy.time_moc_from_json_file(path)
            return cls(IntervalSet(intervals, make_consistent=False), make_consistent=False)
        else:
            formats = ('fits', 'ascii', 'json')
            raise ValueError('format should be one of %s' % (str(formats)))

    def to_string(self, format='ascii'):
        """
        Writes the Time MOC into a string.

        Format can be 'ascii' or 'json', though the json format is not officially supported by the IVOA.

        Parameters
        ----------
        format : str, optional
            The format in which the MOC will be serialized before being saved.
            Possible formats are "ascii" or "json".
            By default, ``format`` is set to "ascii".
        """
        if format == 'ascii':
            return mocpy.time_moc_to_ascii_str(self.max_order, self._interval_set._intervals)
        elif format == 'json':
            return mocpy.time_moc_to_json_str(self.max_order, self._interval_set._intervals)
        else:
            formats = ('ascii', 'json')
            raise ValueError('format should be one of %s' % (str(formats)))

    @classmethod
    def from_string(cls, value, format='ascii'):
        """
        Deserialize the Time MOC from the given string.

        Format can be 'ascii' or 'json', though the json format is not officially supported by the IVOA.

        WARNING: the serialization must be strict, i.e. **must not** contain overlapping elements

        Parameters
        ----------
        format : str, optional
            The format in which the MOC will be serialized before being saved.
            Possible formats are "ascii" or "json".
            By default, ``format`` is set to "ascii".
        """
        if format == 'ascii':
            intervals = mocpy.time_moc_from_ascii_str(value)
            return cls(IntervalSet(intervals, make_consistent=False), make_consistent=False)
        elif format == 'json':
            intervals = mocpy.time_moc_from_json_str(value)
            return cls(IntervalSet(intervals, make_consistent=False), make_consistent=False)
        else:
            formats = ('ascii', 'json')
            raise ValueError('format should be one of %s' % (str(formats)))


"""
"""
from __future__ import absolute_import

import sys
import weakref

import numpy

from .dimensionality import Dimensionality
from . import markup
from .quantity import Quantity, get_conversion_factor
from .registry import unit_registry
from .decorators import memoize, with_doc

__all__ = [
    'CompoundUnit', 'Dimensionless', 'UnitConstant', 'UnitCurrency',
    'UnitCurrent', 'UnitInformation', 'UnitLength', 'UnitLuminousIntensity',
    'UnitMass', 'UnitMass', 'UnitQuantity', 'UnitSubstance', 'UnitTemperature',
    'UnitTime', 'set_default_units'
]

class UnitQuantity(Quantity):

    _primary_order = 90
    _secondary_order = 0
    _reference_quantity = None

    __array_priority__ = 20

    def __new__(
        cls, name, definition=None, symbol=None, u_symbol=None,
        aliases=[], doc=None
    ):
        try:
            assert isinstance(name, str)
        except AssertionError:
            raise TypeError('name must be a string, got %s (not unicode)'%name)
        try:
            assert symbol is None or isinstance(symbol, str)
        except AssertionError:
            raise TypeError(
                'symbol must be a string, '
                'got %s (u_symbol can be unicode)'%symbol
            )

        ret = numpy.array(1, dtype='d').view(cls)
        ret.flags.writeable = False

        ret._name = name
        ret._symbol = symbol
        ret._u_symbol = u_symbol
        if doc is not None:
            ret.__doc__ = doc

        if definition is not None:
            if not isinstance(definition, Quantity):
                definition *= dimensionless
            ret._definition = definition
            ret._conv_ref = definition._reference
        else:
            ret._definition = None
            ret._conv_ref = None

        ret._aliases = aliases

        ret._format_order = (ret._primary_order, ret._secondary_order)
        ret.__class__._secondary_order += 1

        return ret

    def __init__(
        self, name, definition=None, symbol=None, u_symbol=None,
        aliases=[], doc=None
    ):
        unit_registry[name] = self
        if symbol:
            unit_registry[symbol] = self
        for alias in aliases:
            unit_registry[alias] = self

    def __array_finalize__(self, obj):
        pass

    def __hash__(self):
        return hash((type(self), self._name))

    @property
    def _reference(self):
        if self._conv_ref is None:
            return self
        else:
            return self._conv_ref

    @property
    def _dimensionality(self):
        return Dimensionality({self:1})

    @property
    def format_order(self):
        return self._format_order

    @property
    def name(self):
        return self._name

    @property
    def definition(self):
        if self._definition is None:
            return self
        else:
            return self._definition

    @property
    def simplified(self):
        return self._reference.simplified

    @property
    def symbol(self):
        if self._symbol:
            return self._symbol
        else:
            return self.name

    @property
    def u_symbol(self):
        if self._u_symbol:
            return self._u_symbol
        else:
            return self.symbol

    @property
    def units(self):
        return self
    @units.setter
    def units(self, units):
        raise AttributeError('can not modify protected units')

    def __repr__(self):
        ref = self._definition
        if ref:
            ref = ', %s * %s'%(str(ref.magnitude), ref.dimensionality.string)
        else:
            ref = ''
        symbol = self._symbol
        symbol = ', %s'%(repr(symbol)) if symbol else ''
        if markup.config.use_unicode:
            u_symbol = self._u_symbol
            u_symbol = ', %s'%(repr(u_symbol)) if u_symbol else ''
        else:
            u_symbol = ''
        return '%s(%s%s%s%s)'%(
            self.__class__.__name__, repr(self.name), ref, symbol, u_symbol
        )

    @with_doc(Quantity.__str__, use_header=False)
    def __str__(self):
        if self.u_symbol != self.name:
            if markup.config.use_unicode:
                s = '1 %s (%s)'%(self.u_symbol, self.name)
            else:
                s = '1 %s (%s)'%(self.symbol, self.name)
        else:
            s = '1 %s'%self.name

        return s

    @with_doc(Quantity.__add__, use_header=False)
    def __add__(self, other):
        return self.view(Quantity).__add__(other)

    @with_doc(Quantity.__radd__, use_header=False)
    def __radd__(self, other):
        try:
            return self.rescale(other.units).__radd__(other)
        except AttributeError:
            return self.view(Quantity).__radd__(other)

    @with_doc(Quantity.__sub__, use_header=False)
    def __sub__(self, other):
        return self.view(Quantity).__sub__(other)

    @with_doc(Quantity.__rsub__, use_header=False)
    def __rsub__(self, other):
        try:
            return self.rescale(other.units).__rsub__(other)
        except AttributeError:
            return self.view(Quantity).__rsub__(other)

    @with_doc(Quantity.__mod__, use_header=False)
    def __mod__(self, other):
        return self.view(Quantity).__mod__(other)

    @with_doc(Quantity.__rsub__, use_header=False)
    def __rmod__(self, other):
        try:
            return self.rescale(other.units).__rmod__(other)
        except AttributeError:
            return self.view(Quantity).__rmod__(other)

    @with_doc(Quantity.__mul__, use_header=False)
    def __mul__(self, other):
        return self.view(Quantity).__mul__(other)

    @with_doc(Quantity.__rmul__, use_header=False)
    def __rmul__(self, other):
        return self.view(Quantity).__rmul__(other)

    @with_doc(Quantity.__truediv__, use_header=False)
    def __truediv__(self, other):
        return self.view(Quantity).__truediv__(other)

    @with_doc(Quantity.__rtruediv__, use_header=False)
    def __rtruediv__(self, other):
        return self.view(Quantity).__rtruediv__(other)

    if sys.version_info[0] < 3:
        @with_doc(Quantity.__div__, use_header=False)
        def __div__(self, other):
            return self.view(Quantity).__div__(other)

        @with_doc(Quantity.__rdiv__, use_header=False)
        def __rdiv__(self, other):
            return self.view(Quantity).__rdiv__(other)

    @with_doc(Quantity.__pow__, use_header=False)
    def __pow__(self, other):
        return self.view(Quantity).__pow__(other)

    @with_doc(Quantity.__rpow__, use_header=False)
    def __rpow__(self, other):
        return self.view(Quantity).__rpow__(other)

    @with_doc(Quantity.__iadd__, use_header=False)
    def __iadd__(self, other):
        raise TypeError('can not modify protected units')

    @with_doc(Quantity.__isub__, use_header=False)
    def __isub__(self, other):
        raise TypeError('can not modify protected units')

    @with_doc(Quantity.__imul__, use_header=False)
    def __imul__(self, other):
        raise TypeError('can not modify protected units')

    @with_doc(Quantity.__itruediv__, use_header=False)
    def __itruediv__(self, other):
        raise TypeError('can not modify protected units')

    if sys.version_info[0] < 3:
        @with_doc(Quantity.__idiv__, use_header=False)
        def __idiv__(self, other):
            raise TypeError('can not modify protected units')

    @with_doc(Quantity.__ipow__, use_header=False)
    def __ipow__(self, other):
        raise TypeError('can not modify protected units')

    def __getstate__(self):
        """
        Return the internal state of the quantity, for pickling
        purposes.

        """
        state = (1, self._format_order)
        return state

    def __setstate__(self, state):
        ver, fo = state
        self._format_order = fo

    def __reduce__(self):
        """
        Return a tuple for pickling a UnitQuantity.
        """
        return (
            type(self),
            (
                self._name,
                self._definition,
                self._symbol,
                self._u_symbol,
                self._aliases,
                self.__doc__
            ),
            self.__getstate__()
        )

    def copy(self):
        return (
            type(self)(
                self._name,
                self._definition,
                self._symbol,
                self._u_symbol,
                self._aliases,
                self.__doc__
                )
            )

unit_registry['UnitQuantity'] = UnitQuantity

class IrreducibleUnit(UnitQuantity):

    _default_unit = None

    def __init__(
        self, name, definition=None, symbol=None, u_symbol=None,
        aliases=[], doc=None
    ):
        super(IrreducibleUnit, self).__init__(
            name, definition, symbol, u_symbol, aliases, doc
        )
        cls = type(self)
        if cls._default_unit is None:
            cls._default_unit = self

    @property
    def simplified(self):
        return self.view(Quantity).rescale(self.get_default_unit())

    @classmethod
    def get_default_unit(cls):
        return cls._default_unit
    @classmethod
    def set_default_unit(cls, unit):
        if unit is None:
            return
        if isinstance(unit, str):
            unit = unit_registry[unit]
        try:
            # check that conversions are possible:
            get_conversion_factor(cls._default_unit, unit)
        except ValueError:
            raise TypeError('default unit must be of same type')
        cls._default_unit = unit

class UnitMass(IrreducibleUnit):

    _primary_order = 1

class UnitLength(IrreducibleUnit):

    _primary_order = 2

class UnitTime(IrreducibleUnit):

    _primary_order = 3

class UnitCurrent(IrreducibleUnit):

    _primary_order = 4

class UnitLuminousIntensity(IrreducibleUnit):

    _primary_order = 5

class UnitSubstance(IrreducibleUnit):

    _primary_order = 6

class UnitTemperature(IrreducibleUnit):

    _primary_order = 7

class UnitInformation(IrreducibleUnit):

    _primary_order = 8

class UnitCurrency(IrreducibleUnit):

    _primary_order = 9

class CompoundUnit(UnitQuantity):

    _primary_order = 99

    def __new__(cls, name):
        return UnitQuantity.__new__(cls, name, unit_registry[name])

    def __init__(self, name):
        # do not register
        return

    @with_doc(UnitQuantity.__add__, use_header=False)
    def __repr__(self):
        return '1 %s'%self.name

    @property
    def name(self):
        if markup.config.use_unicode:
            return '(%s)'%(markup.superscript(self._name))
        else:
            return '(%s)'%self._name

    def __reduce__(self):
        """
        Return a tuple for pickling a UnitQuantity.
        """
        return (
            type(self),
            (self._name, ),
            self.__getstate__()
            )

    def copy(self):
        return type(self)(self._name)

unit_registry['CompoundUnit'] = CompoundUnit

class Dimensionless(UnitQuantity):

    _primary_order = 100

    def __init__(self, name, definition=None):
        self._name = name

        if definition is None:
            definition = self
        self._definition = definition

        self._format_order = (self._primary_order, self._secondary_order)
        self.__class__._secondary_order += 1

        unit_registry[name] = self

    def __reduce__(self):
        """
        Return a tuple for pickling a UnitQuantity.
        """
        return (
            type(self),
            (
                self._name,
            ),
            self.__getstate__()
        )

    @property
    def _dimensionality(self):
        return Dimensionality()

dimensionless = Dimensionless('dimensionless')

class UnitConstant(UnitQuantity):

    _primary_order = 0

    def __init__(
        self, name, definition=None, symbol=None, u_symbol=None,
        aliases=[], doc=None
    ):
        # we dont want to register constants in the unit registry
        return

def set_default_units(
    system=None, currency=None, current=None, information=None, length=None,
    luminous_intensity=None, mass=None, substance=None, temperature=None,
    time=None
):
    """
    Set the default units in which simplified quantities will be
    expressed.

    system sets the unit system, and can be "SI" or "cgs". All other
    keyword arguments will accept either a string or a unit quantity.
    An error will be raised if it is not possible to convert between
    old and new defaults, so it is not possible to set "kg" as the
    default unit for time.

    If both system and individual defaults are given, the system
    defaults will be applied first, followed by the individual ones.
    """
    if system is not None:
        system = system.lower()
        try:
            assert system in ('si', 'cgs')
        except AssertionError:
            raise ValueError('system must be "SI" or "cgs", got "%s"' % system)
        if system == 'si':
            UnitCurrent.set_default_unit('A')
            UnitLength.set_default_unit('m')
            UnitMass.set_default_unit('kg')
        elif system == 'cgs':
            UnitLength.set_default_unit('cm')
            UnitMass.set_default_unit('g')
        UnitLuminousIntensity.set_default_unit('cd')
        UnitSubstance.set_default_unit('mol')
        UnitTemperature.set_default_unit('degK')
        UnitTime.set_default_unit('s')

    UnitCurrency.set_default_unit(currency)
    UnitCurrent.set_default_unit(current)
    UnitInformation.set_default_unit(information)
    UnitLength.set_default_unit(length)
    UnitLuminousIntensity.set_default_unit(luminous_intensity)
    UnitMass.set_default_unit(mass)
    UnitSubstance.set_default_unit(substance)
    UnitTemperature.set_default_unit(temperature)
    UnitTime.set_default_unit(time)


""" Module for coadding spectra, and comparing customized and default data reduction
"""

from __future__ import (print_function, absolute_import, division, unicode_literals)

import numpy as np
import glob
import pdb

from matplotlib import pyplot as plt

from astropy.table import Table
from astropy.io import fits
import linetools.spectra.xspectrum1d as xspec

from pypeit.core import coadd1d as coadd1d

def flxwave(fname1, fname2s, xlim=(1900, 1950), ylim=None, norm=True, nseg=0, figsize=(6, 5)):
    """ Plot flux vs wavelength for default and customized data reduction

    Parameters
    ----------
    fname1 : str
      a summ file that is output of the default data reduction
    fname2s : list of str
      summ files that are output ot customized data reductions.
      It is needed that these files have spectra taken with the same optical element and central wavelenght as fname1.
    xlim : float
      displayed wavelength range
    ylim : float
      dislayed flux range
    norm : bool
      normalize fluxes to the same median flux
    nseg : 0,1,2 (for segments A, B, C)
      plot flux for segment nseg

    Returns
    -------

    """

    # default data reduction
    fnm = fname1
    tbl = Table.read(fnm)
    wave1 = tbl['WAVELENGTH'][nseg]
    flx1 = tbl['FLUX'][nseg]
    err1 = tbl['ERROR'][nseg]

    plt.figure(figsize=figsize)
    plt.plot(wave1, flx1, color='blue')
    plt.plot(wave1, err1, '--', color='blue', label='previous')

    # customized data reduction
    colors = ['red', 'orange', 'limegreen']
    for i in range(len(fname2s)):
        fname2 = fname2s[i]
        fnm = fname2
        tbl = Table.read(fnm)
        wave2 = tbl['WAVELENGTH'][nseg]
        flx2 = tbl['FLUX'][nseg]
        err2 = tbl['ERROR'][nseg]

        # normalize to median flux
        if norm:
            sc = np.median(flx1) / np.median(flx2)
        else:
            sc = 1
        print(sc)
        flx2 = flx2 * sc
        err2 = err2 * sc

        # plot
        plt.plot(wave2, flx2, color=colors[i])
        plt.plot(wave2, err2, '--', color=colors[i], label='new ' + str(i+1))

    plt.xlim(xlim)
    if ylim is not None:
        plt.ylim(ylim)
    plt.xlabel('Wavelength')
    plt.ylabel('Flux')
    plt.legend()
    plt.show()

def findsp(fldpth, verbose=True, emax=None):
    """ Find all NUV and FUV spectra

    Parameters
    ----------
    fldpth : str
       path to the files
    emax : int
      read extensions 0, 1, ..., emax-1 of sum.fits files (corresponding to segments A,B,C)

    Returns
    -------
    fuvsp : array of XSpectrum1D objects
      FUV spectra
    nuvsp : array of XSpectrum1D objects
      NUV spectra

    """

    # find all spectra
    sumfiles = glob.glob(fldpth + '*sum.fits')

    fuvsp = []
    nuvsp = []
    for i in range(len(sumfiles)):
        spf = sumfiles[i]
        hdu = fits.open(spf)
        # FUV or NUV
        det = hdu[0].header['DETECTOR']
        hdu.close()

        # read data
        tbl = Table.read(spf)
        jmax = len(tbl)
        if (emax is not None):
            if (emax < len(tbl)):
                jmax = emax
        for j in range(jmax):
            wave = tbl['WAVELENGTH'][j]
            flx = tbl['FLUX'][j]
            err = tbl['ERROR'][j]

            # append xspectrum
            xsp = xspec.XSpectrum1D.from_tuple((wave, flx, err))
            if det == 'NUV':
                nuvsp.append(xsp)
            elif det == 'FUV':
                fuvsp.append(xsp)
            else:
                raise IOError("Detector not well defined in ", spf)

    if verbose:
        print('Number of NUV and FUV spectra: ', len(nuvsp), len(fuvsp))
    nuvsp = np.asarray(nuvsp)
    fuvsp = np.asarray(fuvsp)

    return nuvsp, fuvsp

def snf1(wave,flux,error):
    """ Calculate S/N per NUV resolution element = 3 pixels

    Parameters
    ----------
    wave : list of floats
      wavelength
    flux : list of floats
      flux
    error : list of floats
      error in flux
    Returns
    -------
    sn : list of floats
      S/N per resolution element, for each pixel

    """

    sn = []
    for i in np.arange(len(wave)):
        if flux[i] == 0:
            sn.append(0)
        else:
            s = 0
            n = 0
            for j in [-1,0,1]: # resolution element = 3 pixels
                ij = i+j
                if ij < 0:
                    ij = 0
                if ij >= len(flux):
                    ij = len(flux)-1
                s = s + flux[ij]
                n = n + error[ij]**2
            isn = s/(n**0.5)
            sn.append(isn)
    return sn

def snf2(wave,flux,error, R = 19000.):
    """ Calculate S/N per resolution R (for FUV)

    Parameters
    ----------
    wave : list of floats
      wavelength
    flux : list of floats
      flux
    error : list of floats
      error in flux
    Returns
    -------
    sn : list of floats
      S/N per  R = 19000, for each pixel

    """

    sn = []
    for i in np.arange(len(wave)):
        if flux[i] == 0:
            sn.append(0)
        else:
            s = 0
            n = 0
            dlam = wave[i]/R
            lam1 = wave[i] - dlam/2
            lam2 = wave[i] + dlam/2
            pix1 = np.argmin(abs(wave - lam1))
            pix2 = np.argmin(abs(wave - lam2))
            for j in (np.arange(pix2-pix1+1)+pix1):
                s = s + flux[j]
                n = n + error[j]**2
            isn = s/(n**0.5)
            sn.append(isn)
    return sn

def findspsn(spectra,det,minsn=1,verbose=True):
    """ Find all spectra, with median S/N per resolution element greater than minsn.
        Spectra are taken with detector det.

    Parameters
    ----------
    fldpth : str
       path to the files
    det : 'FUV' or 'NUV'
      detector
    minsn : float
      minimum S/N per resolution element, default = 1

    Returns
    -------
    spectrasn : array of XSpectrum1D objects
      spectra with median S/N per resolution element greater than minsn

    """

    spectrasn = []
    for ispec in spectra:
        # S/N
        if det == 'NUV':
            sn = snf1(ispec.wavelength,ispec.flux,ispec.sig)
        elif det == 'FUV':
            sn = snf2(ispec.wavelength,ispec.flux,ispec.sig)
        else:
            raise IOError("det (Detector) could be 'NUV' or 'FUV' only ")
        medsn = np.median(sn)

        # append
        if medsn > minsn:
            spectrasn.append(ispec)
            if verbose:
                print('S/N:',medsn)
        else:
            if verbose:
                print('S/N low ',medsn)

    if verbose:
        print('Number of spectra with S/N > {:f}: {:f}'.format(minsn,len(spectrasn)))

    spectrasn = np.asarray(spectrasn)

    return spectrasn

def medsn(ispec, det):
    """ Find S/N per resolution element for a given spectrum

    Parameters
    ----------
    ispec : XSpectrum1D object
      spectrum
    det : 'NUV' or 'FUV'

    Returns
    -------
    medsn : float
      median S/N per resolution element

    """

    # S/N
    if det == 'NUV':
        sn = snf1(ispec.wavelength, ispec.flux, ispec.sig)
    elif det == 'FUV':
        sn = snf2(ispec.wavelength, ispec.flux, ispec.sig)
    else:
        print('Det options are only NUV and FUV')
        pdb.set_trace()

    medsn = np.median(sn)

    return medsn

def smoothsp(spects, det, snmin, outf=None):
    """ Smooth noisy spectra, and return smoothed spectra with S/N per pixel higher than snmin

    Parameters
    ----------
    spects : list of XSpectrum1D objects
      spectra
    det : 'NUV' or 'FUV'
      'NUV' or 'FUV' detector
    snmin : float
      median S/N per pixel needs to be > snmin

    Returns
    -------
    smspects : list of XSpectrum1D objects
      smoothed spectra with S/N > snmin

    """

    outpspec = []
    print('Total number of spectra: ',len(spects))

    # For each spectrum find median S/N, and boxcar smooth over ism pixels
    # Return only spectra that have S/N > snmin
    for ispec in spects:
        print(ispec)

        # S/N
        imedsn = medsn(ispec, det)
        print(imedsn)

        # sm - lists of the number of pixels to smooth over
        if det == 'NUV':
            sm = [2,3]
        elif det == 'FUV':
            sm = [2,4,8,12,16,20]
            # or define any other list with sm, e.g.:
            # sm = [2,4,6,8,10,12,14,16,18,20]
        else:
            raise IOError("Detector could be 'NUV' or 'FUV' only ")

        # smooth over ism pixels, until the spectrum has S/N > snmin
        if imedsn > snmin:
            outpspec.append(ispec)
        else:
            for ism in sm:
                ispecsm = ispec.box_smooth(ism)
                imedsn = medsn(ispecsm, det)
                print(ism,imedsn)
                if imedsn > snmin:
                    outpspec.append(ispecsm)
                    print('Smoothing with {:f} pixels'.format(ism))
                    break
            if imedsn <= snmin:
                outpspec.append(ispecsm)
                print('Smoothing with {:f} pixels. S/N still low.'.format(ism))

    if outf is not None:
        outpspec[0].write(outf)

    return outpspec

def binsp(spf,kbin = 3, outf=None):
    """ Bin spectrum

        Parameters
        ----------
        spf : file with a XSpectrum1D object
          input spectrum
        kbin : int
          number of pixels to bin
        outf : str
          output file in which the binned spectrum will be written

        Returns
        -------
        xsp2 : XSpectrum1D object
          binned spectrum

        """

    # read spectrum from the file
    xsp = xspec.XSpectrum1D.from_file(spf)
    n = len(xsp.wavelength)
    wave = xsp.wavelength
    flx = xsp.flux
    sig = xsp.sig

    # lists with binned wavelength, flux, error in flux
    waveb = []
    flxb = []
    sigb = []

    # bin
    for j in range(int(n / kbin)):
        i = kbin * j
        waveb.append(np.mean(wave[i:i + kbin].value)) # [angstrom]
        flxb.append(np.mean(flx[i:i + kbin].value))  # [erg /s /cm**2 /angstrom]
        isigb = np.sum((sig[i:i + kbin].value)**2)
        isigb = (isigb**0.5)/kbin
        sigb.append(isigb)

    # output spectrum
    waveb = np.asarray(waveb)
    flxb = np.asarray(flxb)
    sigb = np.asarray(sigb)
    xsp2 = xspec.XSpectrum1D.from_tuple((waveb, flxb, sigb))

    # write
    if outf is not None:
        xsp2.write(outf)

    return xsp2

def coaddspectra(splist,plotsp=True,outf=None,sn_smooth_npix=10):
    """  Coadd spectra

    Parameters
    ----------
    splist : list of XSpectrum1D objects
        List of spectra to coadd
    plotsp : bool
        If True, plot the coadded spectrum
    outf : str
        Output file
    sn_smooth_npix : float
        Parameter in coadd1d.combspec function that defines
        number of pixels to median filter by when computing S/N used to decide how to scale and weight spectra

    Returns
    -------
    sp : XSpectrum1D
        A spectrum that represents coadded spectra from the splist list

    """
    waves = []
    fluxes = []
    ivars = []
    masks = []

    for isp in splist:
        waves.append(isp.wavelength)
        fluxes.append(isp.flux)
        ivars.append(1. / (isp.sig) ** 2.)
        imask = np.repeat(True, len(isp.flux))
        j = np.where((isp.flux == 0) & (isp.sig == 0))[0]
        imask[j] = False
        masks.append(imask)

    waves = np.ndarray.transpose(np.asarray(waves))
    fluxes = np.ndarray.transpose(np.asarray(fluxes))
    ivars = np.ndarray.transpose(np.asarray(ivars))
    masks = np.ndarray.transpose(np.asarray(masks))

    wave_stack, flux_stack, ivar_stack, mask_stack = coadd1d.combspec(
        waves, fluxes, ivars, masks, sn_smooth_npix, show=plotsp)

    ii = np.where(wave_stack > 0)[0]
    coadded_waves = wave_stack[ii]
    coadded_fluxes = flux_stack[ii]
    coadded_sigs = 1 / (np.sqrt(ivar_stack[ii]))

    # write and return the spectrum
    sp = xspec.XSpectrum1D(coadded_waves, coadded_fluxes, coadded_sigs)

    if outf is not None:
        sp.write_to_fits(outf)

    return sp