"""The main container for waveform objects with mode weights"""

import re
import numpy as np
import quaternionic
import spherical
from .. import TimeSeries
from . import WaveformMixin
from .mode_utilities import _expectation_value_LL, _expectation_value_Ldt

NRAR_mode_regex = re.compile(r"""Y_l(?P<L>[0-9]+)_m(?P<M>[-+0-9]+)\.dat""")

class WaveformModes(WaveformMixin, TimeSeries):
    """Array-like object representing time-series data of SWSH modes

    We generally assume that the data represent mode weights in expansions of
    functions in terms of spin-weighted spherical harmonics (where the standard
    spherical harmonics just happen to have spin weight 0).

    This object is based on the TimeSeries object, but has many additional methods
    for manipulating modes.

    Parameters
    ----------
    input_array : (..., N, ..., M, ...) array_like
        Input data representing the dependent variable, in any form that can be
        converted to a numpy array.  This includes scalars, lists, lists of tuples,
        tuples, tuples of tuples, tuples of lists, and numpy ndarrays.  It can have
        an arbitrary number of dimensions, but the length `N` along `time_axis`
        must match the length of `time`, and the length `M` along `modes_axis` (see
        below) must be consistent with `ell_min` and `ell_max`.  Values must be
        finite.
    time : (N,) array_like
        1-D array containing values of the independent variable.  Values must be
        real, finite, and in strictly increasing order.
    time_axis : int, optional
        Axis along which `input_array` is assumed to be varying in time, meaning
        that for `time[i]` the corresponding values are `np.take(input_array, i,
        axis=time_axis)`.  If this is not given, the first axis of `input_array`
        that has the same length as `time` is chosen as the time axis — which may
        be prone to errors.
    modes_axis : int
        Axis of the array along which the modes are stored.  See Notes below.
    ell_min : int
        Smallest value of ℓ stored in the data
    ell_max : int
        Largest value of ℓ stored in the data

    Notes
    -----
    We assume that the modes at a given time, say, are stored in a flat array, in
    order of increasing `m` and `ell`, with `m` varying most rapidly — something
    like the following:

        [f(ell, m) for ell in range(ell_min, ell_max+1) for m in range(-ell,ell+1)]

    The total size is implicitly `ell_max * (ell_max + 2) - ell_min ** 2 + 1`.

    For backwards compatibility, it is possible to retrieve individual modes in the
    same way as the old NRAR-format HDF5 files would be read, as in

        h_22 = waveform["Y_l2_m2.dat"]

    Note that "History.txt" may not contain anything but an empty string, because
    history is not retained in more recent data formats.  Also note that — while
    not strictly a part of this class — the loaders that open waveform files will
    return a dict-like object when the extrapolation order is not specified.  That
    object can also be accessed in a backwards-compatible way much like the root
    directory of the NRAR-format HDF5 files.  For example:

        with sxs.loadcontext("rhOverM_Asymptotic_GeometricUnits_CoM.h5") as f:
            h_22 = f["Extrapolated_N2.dir/Y_l2_m2.dat"]

    This code is identical to the equivalent code using `h5py` except that the call
    to `h5py.File` is replaced with the call to `sxs.loadcontext`.  The `.dat`
    datasets are reconstructed on the fly, but should be bitwise-identical to the
    output from the HDF5 file whenever the underlying format is NRAR.

    """
    import functools

    def __new__(cls, input_array, *args, **kwargs):
        for requirement in ["modes_axis", "ell_min", "ell_max"]:
            if requirement not in kwargs and requirement not in getattr(input_array, "_metadata", {}):
                raise ValueError(f"{cls} could not find required argument '{requirement}'")
        self = super().__new__(cls, input_array, *args, **kwargs)
        return self

    def __getitem__(self, key):
        if isinstance(key, str):
            if key == "History.txt":
                return self._metadata.get("history", "")
            # Assume we're asking for Y_l2_m2.dat or something
            if self.ndarray.shape != (self.n_times, self.n_modes):
                raise ValueError(f"Data has shape {self.ndarray.shape}, which is incompatible with NRAR format")
            match = NRAR_mode_regex.match(key)
            if not match:
                raise ValueError(f"Key '{key}' did not match mode format")
            ell, m = int(match["L"]), int(match["M"])
            if ell < self.ell_min or ell > self.ell_max:
                raise ValueError(f"Key '{key}' requested ell={ell} value not found in this data")
            if abs(m) > ell:
                raise ValueError(f"Key '{key}' requested (ell, m)=({ell}, {m}) value does not make sense")
            index = self.index(ell, m)
            data = np.take(self.ndarray, index, axis=self.modes_axis).view(float).reshape((-1, 2))
            return np.hstack((self.time[:, np.newaxis], data))
        obj, time_key = self._slice(key)
        if time_key is None:
            raise ValueError(f"Fancy indexing (with {key}) is not supported")
        obj = obj.view(type(self))
        if "frame" in obj._metadata and obj.frame.shape == (self.n_times, 4):
            obj._metadata["frame"] = obj.frame[time_key, :]
        return obj

    @property
    def modes_axis(self):
        """Axis of the array storing the various modes

        See the documentation of this class for an explanation of what this means.

        See Also
        --------
        time_axis : Axis of the array along which time varies

        """
        return self._metadata["modes_axis"]

    @property
    def ell_min(self):
        """Smallest value of ℓ stored in the data"""
        return self._metadata["ell_min"]

    @property
    def ell_max(self):
        """Largest value of ℓ stored in the data"""
        return self._metadata["ell_max"]

    @property
    def n_modes(self):
        """Total number of mode weights stored in the data"""
        return self.shape[self.modes_axis]

    @property
    def LM(self):
        """Array of (ell, m) values in the data

        This array is just a flat array of `[ell, m]` pairs.  It is automatically
        recomputed each time `ell_min` or `ell_max` is changed.  Specifically, it is

            np.array([
                [ell, m]
                for ell in range(self.ell_min, self.ell_max+1)
                for m in range(-ell, ell+1)
            ])

        """
        return np.array([[ell, m] for ell in range(self.ell_min, self.ell_max+1) for m in range(-ell, ell+1)])

    def index(self, ell, m):
        """Mode index of given (ell,m) mode in the data

        Parameters
        ----------
        ell : int
        m : int

        Returns
        -------
        idx : int
            Index such that self.LM[idx] == [ell, m]

        """
        return spherical.Yindex(ell, m, self.ell_min)

    @property
    def abs(self):
        """Absolute value of the data

        Returns
        -------
        absolute : TimeSeries
            Because the absolute values make no sense as mode weights, this is just a
            plain TimeSeries object.

        See Also
        --------
        arg

        """
        return np.abs(self.view(TimeSeries))

    @property
    def arg(self):
        """Complex phase angle of the data

        Note that the result is not "unwrapped", meaning that there may be
        discontinuities as the phase approaches ±π.

        Returns
        -------
        phase : TimeSeries
            Values are in the interval (-π, π].

        See Also
        --------
        numpy.angle
        arg_unwrapped

        """
        return np.angle(self.view(TimeSeries))

    @property
    def arg_unwrapped(self):
        """Complex phase angle of the data, unwrapped along the time axis

        The result is "unwrapped", meaning that discontinuities as the phase approaches
        ±π are removed by adding an appropriate amount to all following data points.

        Returns
        -------
        phase : TimeSeries
            Values at the initial time are in the interval (-π, π], but may evolve to
            arbitrary real values.

        See Also
        --------
        numpy.angle
        numpy.unwrap
        arg

        """
        return TimeSeries(np.unwrap(self.arg, axis=self.time_axis), self.time)

    @property
    def norm(self):
        """Compute the L² norm of the waveform

        Returns
        -------
        n : TimeSeries

        See Also
        --------
        numpy.linalg.norm
        numpy.take
        norm2 : squared version of this

        Notes
        -----
        The integral of the (squared) magnitude of the data equals the sum of the
        (squared) magnitude of the modes for orthonormal basis functions, meaning that
        the L² norm of the function equals the basic Euclidean norm of the modes.  We
        assume that these modes are expanded in a band-limited but otherwise complete
        orthonormal basis.

        """
        return TimeSeries(np.linalg.norm(self, axis=self.modes_axis), self.time)

    @property
    def bar(self):
        """Return waveform modes of function representing conjugate of this function

        N.B.: This property is different from the `.conjugate` method; see below.

        See Also
        --------
        re : Return modes of function representing the real part of this function
        im : Return modes of function representing the imaginary part of this function

        Notes
        -----
        This property is different from the `.conjugate` (or `.conj`) method, in that
        `.conjugate` returns the conjugate of the mode weights of the function, whereas
        this property returns the mode weights of the conjugate of the function.  That
        is, `.conjugate` treats the data as a generic numpy array, and simply returns
        the complex conjugate of the raw data without considering what the data
        actually represents.  This property treats the data as a function represented
        by its mode weights.

        The resulting function has the negative spin weight of the input function.

        We have

            conjugate(f){s, l, m} = (-1)**(s+m) * conjugate(f{-s, l, -m})

        """
        return spherical.modes.algebra.bar(self)

    @property
    def re(self):
        """Return waveform modes of function representing real part of this function

        N.B.: This property is different from the `.real` method; see below.

        See Also
        --------
        im : Equivalent method for the imaginary part
        bar : Return modes of function representing the conjugate of this function

        Notes
        -----
        This property is different from the `.real` method, in that `.real` returns the
        real part of the mode weights of the function, whereas this property returns
        the mode weights of the real part of the function.  That is, `.real` treats the
        data as a generic numpy array, and simply returns the real part of the raw data
        without considering what the data actually represents.  This property treats
        the data as a function represented by its mode weights.

        Note that this only makes sense for functions of spin weight zero; taking the
        real part of functions with nonzero spin weight will depend too sensitively on
        the orientation of the coordinate system to make sense.  Therefore, this
        property raises a ValueError for other spins.

        The condition that a function `f` be real is that its modes satisfy

            f{l, m} = conjugate(f){l, m} = (-1)**(m) * conjugate(f{l, -m})

        [Note that conjugate(f){l, m} != conjugate(f{l, m}).]  As usual, we enforce
        that condition by essentially averaging the two modes:

            f{l, m} = (f{l, m} + (-1)**m * conjugate(f{l, -m})) / 2

        """
        return spherical.modes.algebra._real_func(self, False)

    @property
    def im(self):
        """Return waveform modes of function representing imaginary part of this function

        N.B.: This property is different from the `.imag` method; see below.

        See Also
        --------
        re : Equivalent method for the real part
        bar : Return modes of function representing the conjugate of this function

        Notes
        -----
        This property is different from the `.imag` method, in that `.imag` returns the
        imaginary part of the mode weights of the function, whereas this property
        returns the mode weights of the imaginary part of the function.  That is,
        `.imag` treats the data as a generic numpy array, and simply returns the
        imaginary part of the raw data without considering what the data actually
        represents.  This property treats the data as a function represented by its
        mode weights.

        Note that this only makes sense for functions of spin weight zero; taking the
        imaginary part of functions with nonzero spin weight will depend too
        sensitively on the orientation of the coordinate system to make sense.
        Therefore, this property raises a ValueError for other spins.

        The condition that a function `f` be imaginary is that its modes satisfy

            f{l, m} = -conjugate(f){l, m} = (-1)**(m+1) * conjugate(f{l, -m})

        [Note that conjugate(f){l, m} != conjugate(f{l, m}).]  As usual, we enforce
        that condition by essentially averaging the two modes:

            f{l, m} = (f{l, m} + (-1)**(m+1) * conjugate(f{l, -m})) / 2

        """
        return spherical.modes.algebra._imag_func(self, False)

    from spherical.modes.derivatives import (
        Lsquared, Lz, Lplus, Lminus,
        Rsquared, Rz, Rplus, Rminus,
        eth, ethbar
    )

    @property
    def eth_GHP(self):
        """Spin-raising derivative operator defined by Geroch-Held-Penrose

        The operator ð is defined in https://dx.doi.org/10.1063/1.1666410

        See Also
        --------
        eth : Related operator in the Newman-Penrose convention
        ethbar : Similar operator in the Newman-Penrose convention
        ethbar_GHP : Conjugate of this operator

        Notes
        -----
        We assume that the Ricci rotation coefficients satisfy β=β'=0, meaning that
        this operator equals the Newman-Penrose operator ð multiplied by 1/√2.

        """
        return self.eth / np.sqrt(2)

    @property
    def ethbar_GHP(self):
        """Spin-lowering derivative operator defined by Geroch-Held-Penrose

        The operator ð̄ is defined in https://dx.doi.org/10.1063/1.1666410

        See Also
        --------
        eth : Related operator in the Newman-Penrose convention
        ethbar : Similar operator in the Newman-Penrose convention
        eth_GHP : Conjugate of this operator

        Notes
        -----
        We assume that the Ricci rotation coefficients satisfy β=β'=0, meaning that
        this operator equals the Newman-Penrose operator ð̄ multiplied by 1/√2.

        """
        return self.ethbar / np.sqrt(2)

    def max_norm_index(self, skip_fraction_of_data=4):
        """Index of time step with largest norm

        The optional argument skips a fraction of the data.  The default is 4, which
        means that it skips the first 1/4 of the data, and only searches the last 3/4
        of the data for the max.  This must be strictly greater than 1, or the entire
        data is searched for the maximum of the norm.

        """
        if skip_fraction_of_data <= 1:
            return np.argmax(self.norm)
        else:
            i = int(self.n_times // skip_fraction_of_data)
            return np.argmax(self[i:].norm) + i

    def max_norm_time(self, skip_fraction_of_data=4):
        """Return time at which largest norm occurs in data

        See `help(max_norm_index)` for explanation of the optional argument.

        """
        return self.t[self.max_norm_index(skip_fraction_of_data=skip_fraction_of_data)]

    def interpolate(self, new_time, derivative_order=0, out=None):
        """Interpolate this object to a new set of times

        Note that if this object has "frame" data and the derivative order is nonzero,
        it is not entirely clear what is desired.  In those cases, the frame is just
        interpolated to the new times, but no derivative or antiderivative is taken.

        Parameters
        ----------
        new_time : array_like
            Points to evaluate the interpolant at
        derivative_order : int, optional
            Order of derivative to evaluate.  If negative, the antiderivative is
            returned.  Default value of 0 returns the interpolated data without
            derivatives or antiderivatives.  Must be between -3 and 3, inclusive.

        See Also
        --------
        scipy.interpolate.CubicSpline :
            The function that this function is based on.
        antiderivative :
            Calls this funtion with `new_time=self.time` and
            `derivative_order=-antiderivative_order` (defaulting to a single
            antiderivative).
        derivative :
            Calls this function `new_time=self.time` and
            `derivative_order=derivative_order` (defaulting to a single derivative).
        dot :
            Property calling `self.derivative(1)`.
        ddot :
            Property calling `self.derivative(2)`.
        int :
            Property calling `self.antiderivative(1)`.
        iint :
            Property calling `self.antiderivative(2)`.

        """
        result = TimeSeries.interpolate(self, new_time, derivative_order=derivative_order, out=out)
        if self.frame.shape == (self.n_times, 4) and not np.array_equal(self.time, result.time):
            self._metadata["frame"] = quaternionic.squad(self.frame, self.time, result.time)
        return result

    def truncate(self, tol=1e-10):
        """Truncate the precision of this object's data in place

        This function sets bits in the data to 0 when they have lower significance than
        will alter the norm of the Waveform by a fraction `tol` at that instant in
        time.

        """
        if tol != 0.0:
            tol_per_mode = tol / np.sqrt(self.n_modes)
            abs_tolerance = np.linalg.norm(self.ndarray, axis=self.modes_axis, keepdims=True) * tol_per_mode
            super().truncate(abs_tolerance)

    def convert_to_conjugate_pairs(self):
        """Convert modes to conjugate-pair format in place

        This function alters this object's modes to store the sum and difference of
        pairs with opposite `m` values.  If we denote the modes `f[l, m]`, then we
        define

            s[l, m] = (f[l, m] + f̄[l, -m]) / √2
            d[l, m] = (f[l, m] - f̄[l, -m]) / √2

        For m<0 we replace the mode data with `d[l, -m]`, for m=0 we do nothing, and
        for m>0 we replace the mode data with `s[l, m]`.  That is, the mode data on
        output look like this:

            [d[2, 2], d[2, 1], f[2, 0], s[2, 1], s[2, 2], d[3, 3], d[3, 2], ...]

        The factor of √2 is chosen so that the norm (sum of the magnitudes squared) at
        each time for this data is the same as it is for the original data.

        """
        mode_plus = np.empty_like(self.ndarray[..., 0])
        mode_minus = np.empty_like(mode_plus)
        for ell in range(self.ell_min, self.ell_max + 1):
            for m in range(1, ell + 1):
                i_plus = self.index(ell, m)
                i_minus = self.index(ell, -m)
                mode_plus[:] = self.ndarray[..., i_plus]
                mode_minus[:] = self.ndarray[..., i_minus]
                self.ndarray[..., i_plus] = (mode_plus + np.conjugate(mode_minus)) / np.sqrt(2)
                self.ndarray[..., i_minus] = (mode_plus - np.conjugate(mode_minus)) / np.sqrt(2)

    def convert_from_conjugate_pairs(self):
        """Convert modes from conjugate-pair format in place

        This function reverses the effects of `convert_to_conjugate_pairs`.  See that
        function's docstring for details.

        """
        mode_plus = np.empty_like(self.ndarray[..., 0])
        mode_minus = np.empty_like(mode_plus)
        for ell in range(self.ell_min, self.ell_max + 1):
            for m in range(1, ell + 1):
                i_plus = self.index(ell, m)
                i_minus = self.index(ell, -m)
                mode_plus[:] = self.ndarray[..., i_plus]
                mode_minus[:] = self.ndarray[..., i_minus]
                self.ndarray[..., i_plus] = (mode_plus + mode_minus) / np.sqrt(2)
                self.ndarray[..., i_minus] = np.conjugate(mode_plus - mode_minus) / np.sqrt(2)

    def evaluate(self, *directions):
        """Evaluate waveform in a particular direction or set of directions

        Parameters
        ----------
        directions : array_like
            Directions of the observer relative to the source may be specified using
            the usual spherical coordinates, and an optional polarization angle (see
            Notes below).  These can be expressed as 2 or 3 floats (where the third is
            the polarization angle), or as an array with final dimension of size 2 or
            3.  Alternatively, the input may be a `quaternionic.array` (see Notes and
            arxiv.org/abs/1604.08140).  Input arrays can have multiple leading
            dimensions; the final dimension is always considered to hold the
            directions, and the other dimensions are retained in the output.

        Returns
        -------
        signal : array_like
            Note that this is complex-valued, meaning that it represents both
            polarizations.  To get the signal measured by a single detector, just take
            the real part.

        Notes
        -----
        To evaluate mode weights and obtain values, we need to evaluate spin-weighted
        spherical harmonics (SWSHs).  Though usually treated as functions of just the
        angles (θ, ϕ), a mathematically correct treatment (arxiv.org/abs/1604.08140)
        defines SWSHs as functions of the rotation needed to rotate the basis (x̂, ŷ, ẑ)
        onto the usual spherical-coordinate basis (θ̂, ϕ̂, n̂).  This function can take
        quaternionic arrays representing such a rotation directly, or the (θ, ϕ)
        coordinates themselves, and optionally the polarization angle ψ, which are
        automatically converted to quaternionic arrays.

        We define the spherical coordinates (θ, ϕ) such that θ is the polar angle
        (angle between the z axis and the point) and ϕ is the azimuthal angle (angle
        between x axis and orthogonal projection of the point into the x-y plane).
        This gives rise to the standard unit tangent vectors (θ̂, ϕ̂).

        We also define the polarization angle ψ as the angle through which we must
        rotate the vector θ̂ in a positive sense about n̂ to line up with the vector
        defining the legs of the detector.  If not given, this angle is treated as 0.

        Examples
        --------
        We can evaluate the signal in a single direction:

        >>> θ, ϕ, ψ = 0.1, 0.2, 0.3
        >>> w.evaluate(θ, ϕ)  # Default polarization angle
        >>> w.evaluate(θ, ϕ, ψ)  # Specified polarization angle

        Or we can evaluate in a set of directions:

        >>> w.evaluate([[θ, ϕ], [θ+0.4, ϕ], [θ+0.8, ϕ]])

        We can also evaluate on a more extensive set of directions.  Here, we construct
        an equi-angular grid to evaluate the waveform on (though irregular grids are
        also acceptable as long as you can pack them into a numpy array).

        >>> n_theta = n_phi = 2 * w.ell_max + 1
        >>> equiangular = np.array([
            [
                [theta, phi]
                for phi in np.linspace(0.0, 2*np.pi, num=n_phi, endpoint=False)
            ]
            for theta in np.linspace(0.0, np.pi, num=n_theta, endpoint=True)
        ])
        >>> w.evaluate(equiangular)

        """
        if len(directions) == 1:
            directions = directions[0]
        if isinstance(directions, quaternionic.array):
            R = directions
        else:
            directions = np.asarray(directions, dtype=float)
            if directions.shape[-1] == 2:
                R = quaternionic.array.from_spherical_coordinates(directions)
            elif directions.shape[-1] == 3:
                R = quaternionic.array.from_euler_angles(directions[..., 1], directions[..., 0], directions[..., 2])
            else:
                raise ValueError(
                    f"Input `directions` array must be quaternionic, or "
                    f"final dimension of must have size 2 or 3, not {directions.shape[-1]}"
                )

        # Compute the shape of the output
        modes_axis = self.modes_axis
        out_shape = (
            self.shape[:modes_axis]
            + R.shape[:-1]
            + tuple() if (modes_axis % self.ndim) == (-1 % self.ndim) else self.shape[modes_axis+1:]
        )

        # For now, we'll keep the new dimensions flat
        Rflat = R.ndarray.reshape(-1, 4)
        signal_shape = list(self.shape)
        signal_shape[modes_axis] = Rflat.shape[0]
        signal = np.zeros(tuple(signal_shape), dtype=complex)

        # Now, loop through, evaluating for each input R value
        wigner = spherical.Wigner(self.ell_max, ell_min=self.ell_min, mp_max=abs(self.spin_weight))
        sYlm = np.empty(wigner.Ysize, dtype=complex)
        slices = [slice(None) for _ in range(signal.ndim)]
        if np.array_equal(self.frame, np.atleast_2d(quaternionic.one)):  # frame is time-independent
            for i_R in range(Rflat.shape[0]):
                slices[modes_axis] = i_R
                wigner.sYlm(self.spin_weight, Rflat[i_R], out=sYlm)
                signal[tuple(slices)] = np.dot(self.ndarray, sYlm)
        else:  # Need to account for time-dependent frame
            data_slices = [slice(None) for _ in range(self.ndarray.ndim)]
            time_axis = self.time_axis
            for i_t in range(self.n_times):
                slices[time_axis] = i_t
                data_slices[time_axis] = i_t
                R_t = self.frame[i_t].inverse * Rflat
                for i_R, R_i in enumerate(Rflat):
                    slices[modes_axis] = i_R
                    wigner.sYlm(self.spin_weight, R_i, out=sYlm)
                    signal[tuple(slices)] = np.dot(self.ndarray[tuple(data_slices)], sYlm)

        return TimeSeries(signal.reshape(out_shape), self.time)

    # TODO:
    # expectation_value_L
    # # Don't bother with inner_product_LL, as it doesn't appear to be used; maybe a more general version?
    # inner_product
    # mode_frame : Minimally rotating O'Shaughnessy et al. frame
    # to_mode_frame
    # corotating_frame

    @property
    def expectation_value_LL(self):
        """Compute the matrix expectation value ⟨w|LᵃLᵇ|w⟩

        Here, Lᵃ is the usual angular-momentum operator familiar from quantum physics,
        and

            ⟨w|LᵃLᵇ|w⟩ = ℜ{Σₗₘₙ w̄ˡᵐ ⟨l,m|LᵃLᵇ|l,n⟩ wˡⁿ}

        This quantity is important for computing the angular velocity of a waveform.

        See Also
        --------
        expectation_value_Ldt
        angular_velocity

        """
        mode_weights = np.moveaxis(self.ndarray, self.modes_axis, -1)
        output_shape = mode_weights.shape[:-1] + (3, 3)
        mode_weights = mode_weights.reshape(-1, mode_weights.shape[-1])
        LL = np.zeros((mode_weights.shape[0], 3, 3), dtype=float)
        _expectation_value_LL(mode_weights, self.LM, LL)
        return LL.reshape(output_shape)

    @property
    def expectation_value_Ldt(self):
        """Compute the matrix expectation value ⟨w|Lᵃ∂ₜ|w⟩

        Here, Lᵃ is the usual angular-momentum operator familiar from quantum physics,
        ∂ₜ is the partial derivative with respect to time, and

            ⟨w|Lᵃ∂ₜ|w⟩ = ℑ{Σₗₘₙ w̄ˡᵐ ⟨l,m|Lᵃ|l,n⟩ ∂ₜwˡⁿ}

        This quantity is important for computing the angular velocity of a waveform.

        See Also
        --------
        expectation_value_LL
        angular_velocity

        """
        mode_weights = np.moveaxis(self.ndarray, self.modes_axis, -1)
        output_shape = mode_weights.shape[:-1] + (3,)
        mode_weights = mode_weights.reshape(-1, mode_weights.shape[-1])
        mode_weights_dot = np.moveaxis(self.dot.ndarray, self.modes_axis, -1)
        mode_weights_dot = mode_weights_dot.reshape(-1, mode_weights_dot.shape[-1])
        Ldt = np.zeros((mode_weights.shape[0], 3), dtype=float)
        _expectation_value_Ldt(mode_weights, mode_weights_dot, self.LM, Ldt)
        return Ldt.reshape(output_shape)

    @property
    def angular_velocity(self):
        """Angular velocity of waveform

        This function calculates the angular velocity of a WaveformModes object from
        its modes — essentially, the angular velocity of the rotating frame in which
        the time dependence of the modes is minimized.  This was introduced in Sec. II
        of "Angular velocity of gravitational radiation and the corotating frame"
        <http://arxiv.org/abs/1302.2919>.

        It can be calculated in terms of the expectation values ⟨w|Lᵃ∂ₜ|w⟩ and
        ⟨w|LᵃLᵇ|w⟩ according to the relation

            ⟨w|LᵇLᵃ|w⟩ ωₐ = -⟨w|Lᵇ∂ₜ|w⟩

        For each set of modes (e.g., at each instant of time), this is a simple linear
        equation in 3 dimensions to be solved for ω.

        See Also
        --------
        expectation_value_LL
        expectation_value_Ldt

        """
        # Calculate the <L∂ₜ> vector and <LL> matrix at each instant
        ldt = self.expectation_value_Ldt
        ll = self.expectation_value_LL

        # Solve ⟨w|LᵇLᵃ|w⟩ ωₐ = -⟨w|Lᵇ∂ₜ|w⟩ for ω
        ω = -np.linalg.solve(ll, ldt)

        return ω

    def boost(self, v⃗, ell_max):
        """Find modes of waveform boosted by velocity v⃗

        Implements Equation (21) of arxiv.org/abs/1509.00862

        Parameters
        ----------
        v⃗ : array_like
            Three-vector representing the velocity of the boosted frame relative to the
            inertial frame, in units where the speed of light is 1
        ell_max : int
            Maximum value of `ell` to use while computing the transformation, and to
            provide in the returned object.  See Notes, below.

        Returns
        -------
        wprime : WaveformModes
            Modes of waveform measured in boosted frame or of modes from boosted source
            measured in original frame.  This should have the same properties as the
            input waveform, except with (1) different time data [see Notes, below], (2)
            a minimum `ell` value of 0 even for spin weight other than 0, and (3) a
            maximum `ell` value of `ell_max`.

        Notes
        -----
        Due to the nature of the transformation, some of the information in the input
        waveform must be discarded, because it corresponds to slices of the output
        waveform that are not completely represented in the input.  Thus, the times of
        the output waveform will not just be the Lorentz-transformed times of the input
        waveform.

        Depending on the magnitude β=|v⃗|, a very large value of `ell_max` may be
        needed.  The dominant factor is the translation that builds up over time:
        `β*T`, where `T` is the largest time found in the waveform.  For example, if
        β*T ≈ 1000M, we might need `ell_max=64` to maintain a comparable accuracy as in
        the input data.

        Because of the `β*T` effects, it is usually best to set t=0 at the merger time
        — best approximated as `self.max_norm_time()`.  The largest translation is then
        found early in the waveform, when the waveform is changing slowly.

        """
        from .transformations import boost
        return boost(self, v⃗, ell_max)

    def rotate(self, quat):
        """Rotate decomposition basis of modes represented by this waveform

        This returns a new waveform object, with missing "frame" data.

        Parameters
        ----------
        quat : quaternionic.array
            This must have one quaternion or the same number of quaternions as the
            number of times in the waveform.

        """
        from spherical.wigner import _rotate

        R = quaternionic.array(quat)
        wigner = spherical.Wigner(self.ell_max, ell_min=self.ell_min)  #, mp_max=abs(self.spin_weight))
        D = np.zeros(wigner.Dsize, dtype=complex)
        mode_weights = self.ndarray
        rotated_mode_weights = np.zeros_like(mode_weights)
        mode_weights = np.moveaxis(mode_weights, self.modes_axis, -1)
        rotated_mode_weights = np.moveaxis(rotated_mode_weights, self.modes_axis, -1)
        shape = rotated_mode_weights.shape
        if quat.shape == (4,) or quat.shape == (1, 4):
            wigner.D(R, out=D)
            mode_weights = mode_weights.reshape(-1, mode_weights.shape[-1])
            rotated_mode_weights = rotated_mode_weights.reshape(-1, mode_weights.shape[-1])
            _rotate(
                mode_weights, rotated_mode_weights,
                wigner.ell_min, wigner.ell_max, wigner.mp_max,
                self.ell_min, self.ell_max, self.spin_weight,
                D
            )
        elif quat.shape == (self.n_times, 4):
            slices = [slice(None) for _ in range(self.ndim)]
            time_axis = self.time_axis if self.time_axis < self.modes_axis else self.time_axis - 1
            for i_t in range(self.n_times):
                wigner.D(R[i_t], out=D)
                slices[time_axis] = i_t
                s = tuple(slices)
                m = mode_weights[s]
                r = rotated_mode_weights[s]
                m = m.reshape(-1, m.shape[-1])
                r = r.reshape(-1, r.shape[-1])
                _rotate(
                    m, r,
                    wigner.ell_min, wigner.ell_max, wigner.mp_max,
                    self.ell_min, self.ell_max, self.spin_weight,
                    D
                )
        else:
            raise ValueError(
                f"Quaternionic array shape {R.shape} not understood; expected {(4,)}, {(1, 4)}, or {(self.n_times, 4)}"
            )
        rotated_mode_weights = rotated_mode_weights.reshape(shape)
        rotated_mode_weights = np.moveaxis(rotated_mode_weights, -1, self.modes_axis)
        new_metadata = self._metadata.copy()
        new_metadata.pop("frame", None)
        return type(self)(rotated_mode_weights, **new_metadata)

    def to_inertial_frame(self):
        """Return a copy of this waveform in the inertial frame"""
        if "frame" not in self._metadata:
            raise ValueError("This waveform has no frame information")
        if self.frame.shape[0] == 1:
            raise ValueError("This waveform appears to already be in an inertial frame")
        if self.frame.shape != (self.n_times, 4):
            raise ValueError(f"Frame shape {frame.shape} not understood; expected {(self.n_times, 4)}")
        w = self.rotate(~self.frame)
        w._metadata["frame_type"] = "inertial"
        return w

    def to_corotating_frame(
        self, R0=None, tolerance=1e-12, z_alignment_region=None, return_omega=False, truncate_log_frame=False
    ):
        """Return a copy of this waveform in the corotating frame

        The corotating frame is defined to be a rotating frame for which the (L² norm
        of the) time-dependence of the modes expressed in that frame is minimized.
        This leaves the frame determined only up to an overall rotation.  In this

        Parameters
        ----------
        R0 : quaternionic, optional
            Initial value of frame when integrating angular velocity.  Defaults to the
            identity.
        tolerance : float, optional
            Absolute tolerance used in integration of angular velocity
        z_alignment_region : {None, 2-tuple of floats}, optional
            If not None, the dominant eigenvector of the <LL> matrix is aligned with
            the z axis, averaging over this portion of the data.  The first and second
            elements of the input are considered fractions of the inspiral at which to
            begin and end the average.  For example, (0.1, 0.9) would lead to starting
            10% of the time from the first time step to the max norm time, and ending
            at 90% of that time.
        return_omega : bool, optional
            If True, return a 2-tuple consisting of the waveform in the corotating
            frame (the usual returned object) and the angular-velocity data.  That is
            frequently also needed, so this is just a more efficient way of getting the
            data.
        truncate_log_frame : bool, optional
            If True, set bits of log(frame) with lower significance than `tolerance` to
            zero, and use exp(truncated(log(frame))) to rotate the waveform.  Also
            returns `log_frame` along with the waveform (and optionally `omega`)

        """
        raise NotImplementedError()
        frame, omega = corotating_frame(
            self, R0=R0, tolerance=tolerance, z_alignment_region=z_alignment_region, return_omega=True
        )
        if truncate_log_frame:
            log_frame = np.log(frame).ndarray
            power_of_2 = 2 ** int(-np.floor(np.log2(2 * tolerance)))
            log_frame = np.round(log_frame * power_of_2) / power_of_2
            frame = np.exp(quaternionic.array(log_frame))
        w = self.rotate_decomposition_basis(frame)
        w._metadata["frame_type"] = "corotating"
        w._metadata["frame"] = frame
        if return_omega:
            if truncate_log_frame:
                return (w, omega, log_frame)
            else:
                return (w, omega)
        else:
            if truncate_log_frame:
                return (w, log_frame)
            else:
                return w


"""
Radiate.py

Author: Jordan Mirocha
Affiliation: University of Colorado at Boulder
Created on 2010-10-18.

Description: Just keeping this routine around to look at the finite c
implementation - need to transfer it to RadiationField class

"""

import copy, math
import numpy as np

#try:
#    from mpi4py import MPI
#    rank = MPI.COMM_WORLD.rank
#    size = MPI.COMM_WORLD.size
#except ImportError:
#    print "Module mpi4py not found.  No worries, we'll just run in serial."
rank = 0
size = 1

E_th = [13.6, 24.6, 54.4]

neglible_tau = 1e-12
neglible_column = 10

class Radiate:
    def __init__(self, grid, source, **kwargs):
        self.pf = parse_kwargs(**kwargs)
        self.grid = grid
        self.src = source

        # Classes
        #self.cosm = Cosmology(pf)
        #self.control = ControlSimulation(pf)
        self.coeff = RateCoefficients(pf, rs, g)
        #self.Ts = SpinTemperature(pf)
        #self.HI = Hydrogen(pf)

        # Grid/units/etc
        self.GridDimensions = int(pf["GridDimensions"])
        self.grid = np.arange(self.GridDimensions)
        self.StopTime = pf['StopTime'] * self.pf['TimeUnits']
        self.StartCell = pf['StartRadius'] * self.GridDimensions
        self.R0 = pf['StartRadius'] * pf['LengthUnits']

        # Deal with log-grid, compute dx
        #if pf['LogarithmicGrid']:
        #    self.r = np.logspace(np.log10(self.R0), \
        #        np.log10(pf['LengthUnits']), self.GridDimensions + 1)
        #else:
        #    self.r = np.linspace(self.R0, pf['LengthUnits'], self.GridDimensions + 1)

        #self.dx = np.diff(self.r)
        #self.r = self.r[0:-1]
        #self.CellCrossingTime = self.dx / c

        # For convenience
        self.zeros_tmp = np.zeros(4)

    def EvolvePhotonsAtInfiniteSpeed(self, newdata, t, dt, h):
        """
        Solver for InfiniteSpeedOfLight = 1.
        """

        # Could change with time for accreting black holes (or not yet implemented sources)
        Lbol = []
        for rs in self.rs.all_sources:
            Lbol.append(rs.BolometricLuminosity(t))

        # Loop over cells radially, solve rate equations, update values in data -> newdata
        for cell in self.grid:

            # If this cell belongs to another processor, continue
            if self.pf['ParallelizationMethod'] == 1 and size > 1:
                if cell not in self.solve_arr:
                    continue

            # Update progressbar
            if rank == 0 and self.ProgressBar:
                self.pbar.update(cell * size)

            # Read in densities for this cell
            n_e = self.data["ElectronDensity"][cell]
            n_HI = self.data["HIDensity"][cell]
            n_HII = self.data["HIIDensity"][cell]
            n_HeI = self.data["HeIDensity"][cell]
            n_HeII = self.data["HeIIDensity"][cell]
            n_HeIII = self.data["HeIIIDensity"][cell]

            # Read in ionized fractions for this cell
            x_HI = self.x_HI_arr[cell]
            x_HII = self.x_HII_arr[cell]
            x_HeI = self.x_HeI_arr[cell]
            x_HeII = self.x_HeII_arr[cell]
            x_HeIII = self.x_HeIII_arr[cell]

            # Convenience arrays for column, absorber, and ion densities plus a few others
            ncol = self.ncol_all[cell]   # actually log10(ncol)
            nabs = self.nabs_all[cell]
            nion = self.nion_all[cell]
            n_H = n_HI + n_HII
            n_He = n_HeI + n_HeII + n_HeIII
            n_B = n_H + n_He + n_e

            # Read in temperature and internal energy for this cell
            T = self.data["Temperature"][cell]

            # Read radius
            r = self.r[cell]

            # Retrieve indices used for 3D interpolation
            indices = self.indices_all[cell]

            # Retrieve optical depth to this cell
            tau = self.tau_all[cell]

            # Retrieve path length through this cell
            dx = self.dx[cell]

            # Retrieve coefficients and what not.
            args = [nabs, nion, n_H, n_He, n_e]
            args.extend(self.coeff.ConstructArgs(args, indices, Lbol, r, ncol, T, dx, t, self.z, cell))

            # Unpack so we have everything by name
            nabs, nion, n_H, n_He, n_e, Gamma, gamma, Beta, alpha, \
                k_H, zeta, eta, psi, xi, omega, hubble, compton, Jc, Ji = args

            # Sum ionization + heating over all sources (last axis)
            args[5] = np.sum(Gamma, axis = -1)
            args[6] = np.sum(gamma, axis = -1)
            args[9] = np.sum(k_H, axis = -1)

            ######################################
            ######## Solve Rate Equations ########
            ######################################

            self.solver.set_initial_value(self.q_all[cell], t).set_f_params(args).set_jac_params(args)
            self.solver.integrate(t + dt)
            newxHII, newxHeII, newxHeIII, newE = self.solver.y
            newxHII, newxHeII, newxHeIII, newE = self.ConserveParticleNumber(newxHII, newxHeII, newxHeIII, newE)

            # Convert from internal energy back to temperature
            if self.pf['Isothermal']:
                newT = T
            else:
                newT = newE * 2. / 3. / k_B / n_B

            if self.pf['CosmologicalExpansion']:
                n_H = self.cosm.nH0 * (1. + (self.z - self.dz))**3
                n_He = self.cosm.nHe0 * (1. + (self.z - self.dz))**3
                n_e = newxHII * n_H + newxHeII * n_He + 2.0 * newxHeIII * n_He
                n_B = n_H + n_He + n_e

            # Calculate neutral fractions
            newHI = n_H * (1. - newxHII)
            newHII = n_H * newxHII
            newHeI = n_He * (1. - newxHeII - newxHeIII)
            newHeII = n_He * newxHeII
            newHeIII = n_He * newxHeIII

            # Derive Ts, dTb
            delta = 0
            Ts = self.Ts.Ts(self.z, n_H, n_e, newT, newxHII, np.sum(Jc), np.sum(Ji))
            dTb = self.HI.BrightnessTemperature(self.z, newxHII, delta, Ts)

            # Store data
            newdata = self.StoreData(newdata, cell, newHI, newHII, newHeI, newHeII, newHeIII, newT,
                tau, Gamma, k_H, Beta, alpha, zeta, eta, psi, gamma, xi, omega, hubble, compton,
                Jc, Ji, Ts, dTb)

            ######################################
            ################ DONE ################
            ######################################

            # Adjust timestep based on maximum allowed neutral fraction change
            if self.AdaptiveGlobalStep:
                dtphot[cell] = self.control.ComputePhotonTimestep(tau,
                    [newHI, newHeI, newHeII], [newHII, newHeII, newHeIII],
                    ncol, n_H, n_He, n_e, n_B, Gamma, gamma, Beta, alpha, k_H,
                    zeta, eta, psi, xi, omega, hubble, compton, newT, self.z, dt)

        return newdata, dtphot

    def EvolvePhotonsAtFiniteSpeed(self, newdata, t, dt, h):
        """
        Solver for InfiniteSpeedOfLight = 0.

        PhotonPackage guide:
            pack = [EmissionTime, EmissionTimeInterval, ncolHI, ncolHeI, ncolHeII, Energy]
        """

        # Set up timestep array for use on next cycle
        if self.AdaptiveGlobalStep:
            dtphot = 1.e50 * np.ones_like(self.grid)
        else:
            dtphot = dt

        Lbol = self.rs.BolometricLuminosity(t)

        # Photon packages going from oldest to youngest - will have to create it on first timestep
        if t == 0:
            packs = []
        else:
            packs = list(self.data['PhotonPackages'])

        # Add one for this timestep
        packs.append(np.array([t, dt, neglible_column, neglible_column, neglible_column, Lbol * dt]))

        # Loop over photon packages, updating values in cells: data -> newdata
        for j, pack in enumerate(packs):
            t_birth = pack[0]
            r_pack = (t - t_birth) * c        # Position of package before evolving photons
            r_max = r_pack + dt * c           # Furthest this package will get this timestep

            # Cells we need to know about - not necessarily integer
            cell_pack = (r_pack - self.R0) * self.GridDimensions / self.pf['LengthUnits']
            cell_pack_max = (r_max - self.R0) * self.GridDimensions / self.pf['LengthUnits'] - 1
            cell_t = t

            Lbol = pack[-1] / pack[1]

            # Advance this photon package as far as it will go on this global timestep
            while cell_pack < cell_pack_max:

                # What cell are we in
                if cell_pack < 0:
                    cell = -1
                else:
                    cell = int(cell_pack)

                if cell >= self.GridDimensions:
                    break

                # Compute dc (like dx but in fractional cell units)
                # Really how far this photon can go in this step
                if cell_pack % 1 == 0:
                    dc = min(cell_pack_max - cell_pack, 1)
                else:
                    dc = min(math.ceil(cell_pack) - cell_pack, cell_pack_max - cell_pack)

                # We really need to evolve this cell until the next photon package arrives, which
                # is probably longer than a cell crossing time unless the global dt is vv small.
                if (len(packs) > 1) and ((j + 1) < len(packs)):
                    subdt = min(dt, packs[j + 1][0] - pack[0])
                else:
                    subdt = dt

                # If photons haven't hit first cell interface yet, evolve in time
                if cell < 0:
                    cell_pack += dc
                    cell_t += subdt
                    continue

                # Current radius in code units
                r = cell_pack * self.pf['LengthUnits'] / self.pf['GridDimensions']

                # These quantities will be different (in general) for each step
                # of the while loop
                n_e = newdata["ElectronDensity"][cell]
                n_HI = newdata["HIDensity"][cell]
                n_HII = newdata["HIIDensity"][cell]
                n_HeI = newdata["HeIDensity"][cell]
                n_HeII = newdata["HeIIDensity"][cell]
                n_HeIII = newdata["HeIIIDensity"][cell]
                n_H = n_HI + n_HII
                n_He = n_HeI + n_HeII + n_HeIII

                # Read in ionized fractions for this cell
                x_HI = n_HI / n_H
                x_HII = n_HII / n_H
                x_HeI = n_HeI / n_He
                x_HeII = n_HeII = n_He
                x_HeIII = n_HeIII = n_He

                # Compute mean molecular weight for this cell
                mu = 1. / (self.cosm.X * (1. + x_HII) + self.cosm.Y * (1. + x_HeII + x_HeIII) / 4.)

                # Retrieve path length through this cell
                dx = self.dx[cell]

                # Crossing time
                dct = self.CellCrossingTime[cell]

                # For convenience
                nabs = np.array([n_HI, n_HeI, n_HeII])
                nion = np.array([n_HII, n_HeII, n_HeIII])
                n_H = n_HI + n_HII
                n_He = n_HeI + n_HeII + n_HeIII
                n_B = n_H + n_He + n_e

                # Compute internal energy for this cell
                T = newdata["Temperature"][cell]
                E = 3. * k_B * T * n_B / mu / 2.

                q_cell = [n_HII, n_HeII, n_HeIII, E]

                # Add columns of this cell
                packs[j][2] += newdata['HIDensity'][cell] * dc * dx
                packs[j][3] += newdata['HeIDensity'][cell] * dc * dx
                packs[j][4] += newdata['HeIIDensity'][cell] * dc * dx
                ncol = np.log10(packs[j][2:5])

                ######################################
                ######## Solve Rate Equations ########
                ######################################

                # Retrieve indices used for interpolation
                indices = self.coeff.Interpolate.GetIndices(ncol)

                # Retrieve coefficients and what not.
                args = [nabs, nion, n_H, n_He, n_e]
                args.extend(self.coeff.ConstructArgs(args, indices, Lbol, r, ncol, T, dx * dc, t, self.z, cell))

                # Unpack so we have everything by name
                nabs, nion, n_H, n_He, n_e, Gamma, gamma, Beta, alpha, \
                    k_H, zeta, eta, psi, xi, omega, hubble, compton = args

                ######################################
                ######## Solve Rate Equations ########
                ######################################

                tarr, qnew, h, odeitr, rootitr = self.solver.integrate(self.RateEquations,
                    q_cell, cell_t, cell_t + subdt, self.z, self.z - self.dz, None, h, *args)

                # Unpack results of coupled equations - Remember: these are lists and we only need the last entry
                newHII, newHeII, newHeIII, newE = qnew

                # Weight by volume if less than a cell traversed
                if dc < 1:
                    dnHII = newHII - newdata['HIIDensity'][cell]
                    dnHeII = newHeII - newdata['HeIIDensity'][cell]
                    dnHeIII = newHeIII - newdata['HeIIIDensity'][cell]
                    dV = self.coeff.ShellVolume(r, dx * dc) / self.coeff.ShellVolume(self.r[cell], dx)
                    newHII = newdata['HIIDensity'][cell] + dnHII * dV
                    newHeII = newdata['HeIIDensity'][cell] + dnHeII * dV
                    newHeIII = newdata['HeIIIDensity'][cell] + dnHeIII * dV

                # Calculate neutral fractions
                newHI = n_H - newHII
                newHeI = n_He - newHeII - newHeIII

                # Convert from internal energy back to temperature
                newT = newE * 2. * mu / 3. / k_B / n_B

                # Store data
                newdata = self.StoreData(newdata, cell, newHI, newHII, newHeI, newHeII, newHeIII, newT,
                    self.tau_all[cell], odeitr, h, rootitr, Gamma, k_H, Beta, alpha, zeta, eta, psi, gamma, xi,
                    omega, hubble, compton)

                cell_pack += dc
                cell_t += subdt

                ######################################
                ################ DONE ################
                ######################################

        # Adjust timestep for next cycle
        if self.AdaptiveGlobalStep:
            n_HI = newdata['HIDensity']
            n_HII = newdata['HIIDensity']
            n_H_all = n_HI + n_HII
            n_HeI = newdata['HeIDensity']
            n_HeII = newdata['HeIIDensity']
            n_HeIII = newdata['HeIIIDensity']
            n_He_all = n_HeI + n_HeII + n_HeIII
            n_e_all = n_HII + n_HeII + 2 * n_HeIII
            T = newdata['Temperature']
            n_B_all = n_H_all + n_He_all + n_e_all

            ncol_HI = np.roll(np.cumsum(n_HI * self.dx), 1)
            ncol_HeI = np.roll(np.cumsum(n_HeI * self.dx), 1)
            ncol_HeII = np.roll(np.cumsum(n_HeII * self.dx), 1)
            ncol_HI[0] = ncol_HeI[0] = ncol_HeII[0] = neglible_column
            ncol = np.transpose(np.log10([ncol_HI, ncol_HeI, ncol_HeII]))

            tau = self.ComputeOpticalDepths([ncol_HI, ncol_HeI, ncol_HeII])

            for cell in self.grid:
                r = self.r[cell]
                dx = self.dx[cell]
                nabs = np.array([n_HI[cell], n_HeI[cell], n_HeII[cell]])
                nion = np.array([n_HII[cell], n_HeII[cell], n_HeIII[cell]])

                indices = self.coeff.Interpolate.GetIndices(ncol[cell])

                args = [nabs, nion, n_H_all[cell], n_He_all[cell], n_e_all[cell]]
                args.extend(self.coeff.ConstructArgs(args, indices, Lbol, r, ncol[cell], T[cell], dx, t, self.z))

                # Unpack so we have everything by name
                nabs, nion, n_H, n_He, n_e, Gamma, gamma, Beta, alpha, \
                    k_H, zeta, eta, psi, xi, omega, hubble, compton = args

                dtphot[cell] = self.control.ComputePhotonTimestep(tau[:,cell],
                    nabs, nion, ncol[cell], n_H, n_He,
                    n_e, n_B_all[cell], Gamma, gamma, Beta, alpha, k_H, zeta,
                    eta, psi, xi, omega, hubble, compton, T[cell], self.z, dt)

                if self.pf['LightCrossingTimeRestrictedTimestep']:
                    dtphot[cell] = min(dtphot[cell],
                        self.LightCrossingTimeRestrictedTimestep * self.CellCrossingTime[cell])

        # Update photon packages
        newdata['PhotonPackages'] = np.array(self.UpdatePhotonPackages(packs, t + dt))

        return newdata, dtphot

    #def ConserveParticleNumber(self, newxHII, newxHeII, newxHeIII, newE):
    #    """
    #    Take solver.y and make sure ion fractions sum to 1 - MinimumSpeciesFraction.
    #    """
    #
    #    if newxHII >= 1:
    #        newxHII = 1. - self.pf['MinimumSpeciesFraction']
    #
    #    if (newxHeII + newxHeIII) >= 1:
    #        if newxHeII > newxHeIII:
    #            newxHeII += 1. - (newxHeII + newxHeIII)
    #            newxHeIII -= self.pf['MinimumSpeciesFraction']
    #        else:
    #            newxHeIII += 1. - (newxHeII + newxHeIII)
    #            newxHeII -= self.pf['MinimumSpeciesFraction']
    #
    #    return newxHII, newxHeII, newxHeIII, newE
    #
    #def UpdatePhotonPackages(self, packs, t_next):
    #    """
    #    Get rid of old photon packages.
    #    """
    #
    #    to_eliminate = []
    #    for i, pack in enumerate(packs):
    #        if (t_next - pack[0]) > (self.pf['LengthUnits'] / c):
    #            to_eliminate.append(i)
    #
    #    to_eliminate.reverse()
    #    for element in to_eliminate:
    #        packs.pop(element)
    #
    #    return np.array(packs)
    #
    #def StoreData(self, newdata, cell, newHI, newHII, newHeI, newHeII, newHeIII, newT,
    #    tau, Gamma, k_H, Beta, alpha, zeta, eta, psi, gamma, xi, omega, hubble, compton,
    #    Jc, Ji, Ts, dTb, packs = None):
    #    """
    #    Copy fields to newdata dictionary.
    #    """
    #
    #    # Update quantities in 'data' -> 'newdata'
    #    newdata["HIDensity"][cell] = newHI
    #    newdata["HIIDensity"][cell] = newHII
    #    newdata["HeIDensity"][cell] = newHeI
    #    newdata["HeIIDensity"][cell] = newHeII
    #    newdata["HeIIIDensity"][cell] = newHeIII
    #    newdata["ElectronDensity"][cell] = newHII + newHeII + 2.0 * newHeIII
    #    newdata["Temperature"][cell] = newT
    #    newdata["OpticalDepth"][cell] = tau
    #    newdata["SpinTemperature"][cell] = Ts
    #    newdata["BrightnessTemperature"][cell] = dTb
    #
    #    if self.pf['OutputRates']:
    #        for i in xrange(3):
    #
    #            for j in xrange(self.rs.Ns):
    #                newdata['PhotoIonizationRate%i_src%i' % (i, j)][cell] = Gamma[i][j]
    #                newdata['PhotoHeatingRate%i_src%i' % (i, j)][cell] = k_H[i][j]
    #                newdata['SecondaryIonizationRate%i_src%i' % (i, j)][cell] = gamma[i,:,j]
    #
    #                newdata['InjectedLyAFlux%i_src%i' % (i, j)][cell] = Ji[i][j]
    #
    #                if i == 0:
    #                    newdata['ContinuumLyAFlux_src%i' % j][cell] = Jc[j]
    #
    #            newdata['CollisionalIonizationRate%i' % i][cell] = Beta[i]
    #            newdata['RadiativeRecombinationRate%i' % i][cell] = alpha[i]
    #            newdata['CollisionalIonizationCoolingRate%i' % i][cell] = zeta[i]
    #            newdata['RecombinationCoolingRate%i' % i][cell] = eta[i]
    #            newdata['CollisionalExcitationCoolingRate%i' % i][cell] = psi[i]
    #
    #            if i == 2:
    #                newdata['DielectricRecombinationRate'][cell] = xi[i]
    #                newdata['DielectricRecombinationCoolingRate'][cell] = omega[i]
    #
    #        newdata['HubbleCoolingRate'][cell] = hubble
    #
    #    return newdata
    #
    #def RateEquations(self, t, q, args):
    #    """
    #    This function returns the right-hand side of our ODE's (Equations 1, 2, 3 and 9 in Mirocha et al. 2012).
    #
    #    q = [x_HII, x_HeII, x_HeIII, E] - our four coupled equations. q = generalized quantity I guess.
    #
    #    args = (nabs, nion, n_H, n_He, n_e, Gamma, gamma, Beta, alpha, k_H, zeta, eta, psi, xi, omega)
    #
    #    where nabs, nion, Gamma, gamma, Beta, alpha, k_H, zeta, eta, psi, xi = 3 element arrays
    #        (one entry per species)
    #
    #    """
    #
    #    nabs, nion, n_H, n_He, n_e, Gamma, gamma, Beta, alpha, k_H, zeta, eta, \
    #        psi, xi, omega, hubble, compton, Jc, Ji = args
    #
    #    dqdt = self.zeros_tmp
    #
    #    # Neutrals (current time-step)
    #    nHI = n_H * (1. - q[0])
    #    nHII = n_H * q[0]
    #    nHeI = n_He * (1. - q[1] - q[2])
    #    nHeII = n_He * q[1]
    #    nHeIII = n_He * q[2]
    #
    #    nabs = [nHI, nHeI, nHeII]
    #    nion = [nHII, nHeII, nHeIII]
    #
    #    xHI = 1. - q[0]
    #    xHII = q[0]
    #    xHeI = 1. - q[1] - q[2]
    #    xHeII = q[1]
    #    xHeIII = q[2]
    #
    #    n_e = q[0] * n_H + q[1] * n_He + 2. * q[2] * n_He
    #
    #    # Always solve hydrogen rate equation
    #    dqdt[0] = (Gamma[0] + Beta[0] * n_e) * xHI + \
    #              (gamma[0][0] * xHI + gamma[0][1] * nHeI / n_H + gamma[0][2] * nHeII / n_H) - \
    #               alpha[0] * n_e * xHII
    #
    #    # Helium rate equations
    #    if self.pf['MultiSpecies']:
    #        dqdt[1] = (Gamma[1] + Beta[1] * n_e) * xHeI + \
    #                  (gamma[1][0] * nHI / n_He + gamma[1][1] * nHeI / n_He + gamma[1][2] * xHeII / n_He)  + \
    #                   alpha[2] * n_e * xHeIII - \
    #                  (Beta[1] + alpha[1] + xi[1]) * n_e * nHeII / n_He
    #
    #        dqdt[2] = (Gamma[2] + Beta[2] * n_e) * xHeII - alpha[2] * n_e * nHeIII / n_He
    #
    #    # Temperature evolution - looks dumb but using np.sum is slow
    #    if not self.pf['Isothermal']:
    #        phoheat = k_H[0] * nHI
    #        ioncool = zeta[0] * nHI
    #        reccool = eta[0] * nHII
    #        exccool = psi[0] * nHI
    #
    #        if self.pf['MultiSpecies']:
    #            phoheat += k_H[1] * nabs[1] + k_H[2] * nabs[2]
    #            ioncool += zeta[1] * nabs[1] + zeta[2] * nabs[2]
    #            reccool += eta[1] * nion[1] + eta[2] * nion[2]
    #            exccool += psi[1] * nabs[1] + psi[2] * nabs[2]
    #
    #        dqdt[3] = phoheat - n_e * (ioncool + reccool + exccool + nHeIII * omega[1])
    #
    #        if self.pf['CosmologicalExpansion']:
    #            dqdt[3] -= 2. * hubble * q[3]
    #
    #            #if self.pf['ComptonHeating']:
    #            #    dqdt[3] += compton
    #
    #    return dqdt
    #
    #def Jacobian(self, t, q, args):
    #    """
    #    Jacobian of the rate equations.
    #    """
    #
    #    nabs, nion, n_H, n_He, n_e, Gamma, gamma, Beta, alpha, k_H, zeta, eta, \
    #        psi, xi, omega, hubble, compton = args
    #
    #    # Neutrals (current time-step)
    #    nHI = n_H * (1. - q[0])
    #    nHII = n_H * q[0]
    #    nHeI = n_He * (1. - q[1] - q[2])
    #    nHeII = n_He * q[1]
    #    nHeIII = n_He * q[2]
    #
    #    xHI = 1. - q[0]
    #    xHII = q[0]
    #    xHeI = 1. - q[1] - q[2]
    #    xHeII = q[1]
    #    xHeIII = q[2]
    #
    #    J = np.zeros([4, 4])
    #
    #    J[0][0] = -(Gamma[0] + Beta[0] * n_e - gamma[0][0]) - alpha[0] * n_e
    #    J[0][1] = -gamma[0][1] / n_H + gamma[0][2] / n_H
    #    J[0][2] = -gamma[0][1] / n_H
    #    J[1][0] = -gamma[1][0] / n_He
    #    J[1][1] = -(Gamma[1] + Beta[0] * n_e + gamma[1][1] / n_He - gamma[1][2] / n_He) \
    #              +(Beta[1] + alpha[1] + xi[1]) * n_e / n_He
    #    J[1][2] = -gamma[1][1] / n_He + alpha[2] * n_e
    #    J[2][1] = (Gamma[2] + Beta[2] * n_e)
    #    J[2][2] = -alpha[2] * n_e / n_He
    #    J[3][0] = n_e * (zeta[0] - eta[0] + psi[0]) - k_H[0]
    #    J[3][1] = n_e * (zeta[1] - eta[1] + psi[1]) - k_H[1]
    #    J[3][2] = n_e * (zeta[2] - eta[2] + psi[2] - omega[1]) - k_H[2]
    #
    #    return J
    #
    #def ComputeOpticalDepths(self, ncol):
    #    """
    #    Compute optical depths *between* source and all cells.  Used solely (I think)
    #    for calculating next timestep.
    #    """
    #
    #    tau_all_arr = np.zeros([3, self.GridDimensions, self.pf['NumberOfSources']])
    #    for rs, source in enumerate(self.rs.all_sources):
    #        if not source.TableAvailable:
    #            tmp_nHI = np.transpose(len(source.E) * [ncol[0]])
    #            tau_all_arr[0,:,rs] = np.sum(tmp_nHI * source.sigma[0, rs], axis = 1)
    #
    #            if self.pf['MultiSpecies']:
    #                tmp_nHeI = np.transpose(len(source.E) * [ncol[1]])
    #                tmp_nHeII = np.transpose(len(source.E) * [ncol[2]])
    #                tau_all_arr[1,:,rs] = np.sum(tmp_nHeI * source.sigma[1, rs], axis = 1)
    #                tau_all_arr[2,:,rs] = np.sum(tmp_nHeII * source.sigma[2, rs], axis = 1)
    #        else:
    #            for i, col in enumerate(np.log10(ncol[0])):
    #                tau_all_arr[0,i,rs] = source.Interpolate.OpticalDepth(col, 0)
    #
    #                if self.pf['MultiSpecies']:
    #                    tau_all_arr[1,i,rs] = source.Interpolate.OpticalDepth(np.log10(ncol[1][i]), 1)
    #                    tau_all_arr[2,i,rs] = source.Interpolate.OpticalDepth(np.log10(ncol[2][i]), 2)
    #
    #            tau_all_arr[0][0][rs] = tau_all_arr[1][0] = tau_all_arr[2][0] = neglible_tau
    #
    #    # Take minimum optical depth over sources to be conservative with dt
    #    return np.min(tau_all_arr, axis = -1)
    #
#    def EvolvePhotons(self, data, t, dt, h, lb):
    #    """
    #    This routine calls our solvers and updates 'data' -> 'newdata'
    #    """
    #
    #    # Make data globally accessible
    #    self.data = data
    #
    #    # Figure out which processors will solve which cells and create newdata dict
    #    self.solve_arr, newdata = self.control.DistributeDataAcrossProcessors(data, lb)
    #
    #    # If we're in an expanding universe, prepare to dilute densities by (1 + z)**3
    #    self.z = 0
    #    self.dz = 0
    #    if self.pf['CosmologicalExpansion']:
    #        self.z = self.cosm.TimeToRedshiftConverter(0., t, self.pf['InitialRedshift'])
    #        self.dz = dt / self.cosm.dtdz(self.z)
    #
    #    # Nice names for densities, ionized fractions
    #    self.n_H_arr = data["HIDensity"] + data["HIIDensity"]
    #    self.x_HI_arr = data["HIDensity"] / self.n_H_arr
    #    self.x_HII_arr = data["HIIDensity"] / self.n_H_arr
    #
    #    if self.pf['MultiSpecies']:
    #        self.n_He_arr = data["HeIDensity"] + data["HeIIDensity"] + data["HeIIIDensity"]
    #        self.x_HeI_arr = data["HeIDensity"] / self.n_He_arr
    #        self.x_HeII_arr = data["HeIIDensity"] / self.n_He_arr
    #        self.x_HeIII_arr = data["HeIIIDensity"] / self.n_He_arr
    #    else:
    #        self.n_He_arr = self.x_HeI_arr = self.x_HeII_arr = self.x_HeIII_arr = np.zeros_like(self.x_HI_arr)
    #
    #    # Compute column densities - meaning column density *between* source and cell
    #    self.ncol_HI = np.roll(np.cumsum(data["HIDensity"] * self.dx), 1)
    #    self.ncol_HeI = np.roll(np.cumsum(data["HeIDensity"] * self.dx), 1)
    #    self.ncol_HeII = np.roll(np.cumsum(data["HeIIDensity"] * self.dx), 1)
    #    self.ncol_HI[0] = self.ncol_HeI[0] = self.ncol_HeII[0] = neglible_column
    #
    #    # Convenience arrays for column densities, absorbers, ion densities, and some others
    #    self.ncol_all = np.transpose(np.log10([self.ncol_HI, self.ncol_HeI, self.ncol_HeII]))
    #    self.nabs_all = np.transpose([data["HIDensity"], data["HeIDensity"], data["HeIIDensity"]])
    #    self.nion_all = np.transpose([data["HIIDensity"], data["HeIIDensity"], data["HeIIIDensity"]])
    #    self.mu_all = 1. / (self.cosm.X * (1. + self.x_HII_arr) + self.cosm.Y * (1. + self.x_HeII_arr + self.x_HeIII_arr) / 4.)
    #    self.ne_all = data["ElectronDensity"]
    #    self.nB_all = self.n_H_arr + self.n_He_arr + self.ne_all
    #
    #    self.q_all = np.transpose([self.x_HII_arr, self.x_HeII_arr, self.x_HeIII_arr,
    #        3. * k_B * self.nB_all * data["Temperature"] / 2.])
    #
    #    # Retrieve indices used for N-D interpolation
    #    self.indices_all = []
    #    for i, col in enumerate(self.ncol_all):
    #        tmp = []
    #        for rs in self.rs.all_sources:
    #            if rs.TableAvailable:
    #                tmp.append(rs.Interpolate.GetIndices([col[0], col[1], col[2], np.log10(self.x_HII_arr[i]), t]))
    #            else:
    #                tmp.append(None)
    #        self.indices_all.append(tmp)
    #
    #    # Compute tau *between* source and all cells
    #    tau_all_arr = self.ComputeOpticalDepths([self.ncol_HI, self.ncol_HeI, self.ncol_HeII])
    #    self.tau_all = zip(*tau_all_arr)
    #
    #    # Print status, and update progress bar
    #    if rank == 0:
    #        print "rt1d: %g < t < %g" % (t / self.pf['TimeUnits'], (t + dt) / self.pf['TimeUnits'])
    #    if rank == 0 and self.ProgressBar:
    #        self.pbar = ProgressBar(widgets = widget, maxval = self.grid[-1]).start()
    #
    #    # SOLVE: c -> inf
    #    if self.pf['InfiniteSpeedOfLight']:
    #        newdata, dtphot = self.EvolvePhotonsAtInfiniteSpeed(newdata, t, dt, h)
    #
    #    # SOLVE: c = finite
    #    else:
    #        newdata, dtphot = self.EvolvePhotonsAtFiniteSpeed(newdata, t, dt, h)
    #
    #    ###
    #    ## Tidy up a bit
    #    ###
    #
    #    # If multiple processors at work, communicate data and timestep
    #    if (size > 1) and (self.pf['ParallelizationMethod'] == 1):
    #        for key in newdata.keys():
    #            newdata[key] = MPI.COMM_WORLD.allreduce(newdata[key], newdata[key])
    #
    #        dtphot = MPI.COMM_WORLD.allreduce(dtphot, dtphot)
    #
    #    # Compute timestep for next cycle based on minimum dt required over entire grid
    #    if self.AdaptiveGlobalStep:
    #        newdt = min(np.min(dtphot), 2 * dt)
    #
    #        if self.pf['CosmologicalExpansion'] and self.pf['RedshiftRestrictedTimestep']:
    #            newdt = min(newdt, self.cosm.dtdz(self.z) * self.pf['MaxRedshiftStep'])
    #
    #    else:
    #        newdt = dt
    #
    #    # Store timestep information
    #    newdata['dtPhoton'] = dtphot
    #
    #    # Load balance grid for next timestep
    #    if size > 1 and self.pf['ParallelizationMethod'] == 1:
    #        lb = self.control.LoadBalance(dtphot)
    #    else:
    #        lb = None
    #
    #    if rank == 0 and self.ProgressBar:
    #        self.pbar.finish()
    #
    #    return newdata, h, newdt, lb
    #


#!/usr/bin/env python
import os
import sys
import numpy as np
from glob import glob
import matplotlib
matplotlib.use('PDF')
from mgplottools.mpl import set_axis, new_figure, get_color, set_color_cycle, ls

def create_figure(outfile, data):
    """
    data is a dict
    run_folder => {
        'E' => array of amplitudes
        'concurrence' => array of concurrences for each value in E,
        'pop_loss' => array of population loss from log subsp (E)
        'nq_vals' => array of number of qubit levels included in prop (E),
        'nc_vals' => array of number of qubit levels included in prop (E)
    }
    """

    # Layout
    fig_width       = 12.5              # Total canvas (cv) width
    legend_offset   = fig_width-3.9     # Left cv -> legend
    left_margin     = 1.1               # Left cv -> plot area
    right_margin    = 4.1               # plot area -> right cv
    top_margin      = 0.4               # top cv -> plot area
    bottom_margin   = 1.0               # bottom cv -> plot area
    h               = 2.0               # height of each panel
    gap             = 0.0               # gap between panels
    w = fig_width - (left_margin + right_margin)  # width of panel

    n_panels = 2
    panels_total_height = n_panels*h + (n_panels-1)*gap

    fig_height = bottom_margin + n_panels*h + (n_panels-1)*gap + top_margin
    fig = new_figure(fig_width, fig_height)

    E = data[data.keys()[0]]['E'] # pulse amplitude, for all axes

    axes = [] # panels
    # create axes from top to bottom
    for i in xrange(n_panels):
        p_offset = bottom_margin + (n_panels-i-1)*(h+gap)
        pos = [left_margin/fig_width, p_offset/fig_height,
               w/fig_width, h/fig_height]
        ax = fig.add_axes(pos)
        if i == n_panels-1:
            set_axis(ax, 'x', E[0], E[-1], 50, minor=5,
                     label=r'peak pulse amplitude $\epsilon_0$ (MHz)')
        else:
            set_axis(ax, 'x', E[0], E[-1], 50, minor=5, ticklabels=False)
        set_axis(ax, 'y', 0, 1, 0.5, minor=5)
        if i == 0: # top panel
            ax.set_yticklabels(['0', '0.5', '1.0'])
        else:
            ax.set_yticklabels(['0', '0.5', ''])
        axes.append(ax)

    titles = [# title for each panel (top to bottom)
        r'$\omega_c = \SI{6.0}{GHz} < \omega_1, \omega_2$',
        r'$\omega_c = \SI{8.1}{GHz} > \omega_1, \omega_2$',
    ]

    # do the plotting
    for (file, label, color, line_style) in (
        ('params1d-60_T200', r'$\omega_d = \omega_c - 60$~MHz', 'blue',   'solid'),
        ('params1d-40_T200', r'$\omega_d = \omega_c - 40$~MHz', 'orange', 'solid'),
        ('params1d40_T200',  r'$\omega_d = \omega_c + 40$~MHz', 'red',    'long-dashed'),
        ('params1d60_T200',  r'$\omega_d = \omega_c + 60$~MHz', 'green',  'long-dashed')
     ):
        try:
            axes[0].plot(data[file]['E'], data[file]['concurrence'],
                        label=label, color=get_color(color),
                        dashes=ls[line_style])
        except KeyError:
            pass
    for (file, label, color, line_style) in (
        ('params2d-60_T200', r'$\omega_d = \omega_c - 60$~MHz', 'blue',   'solid'),
        ('params2d-40_T200', r'$\omega_d = \omega_c - 40$~MHz', 'orange', 'solid'),
        ('params2d40_T200',  r'$\omega_d = \omega_c + 40$~MHz', 'red',    'long-dashed'),
        ('params2d60_T200',  r'$\omega_d = \omega_c + 60$~MHz', 'green',  'long-dashed')
     ):
        try:
            axes[1].plot(data[file]['E'], data[file]['concurrence'],
                        label=label, color=get_color(color),
                        dashes=ls[line_style])
        except KeyError:
            pass

    # label points
    marker_x = (125, 250)
    marker_y = (0.9984, 0.9857)
    marker_text = ("1", "2")
    axes[0].scatter(marker_x,  marker_y, c='black', clip_on=False,
                    edgecolors='none', zorder=10)
    for i, txt in enumerate(marker_text):
            axes[0].annotate(txt, (marker_x[i],marker_y[i]),
                             xytext=(2, 3), textcoords='offset points')
    marker_x = (425,)
    marker_y = (0.9968, )
    marker_text = ("3",)
    axes[1].scatter(marker_x,  marker_y, c='black', clip_on=False,
                    edgecolors='none', zorder=10)
    for i, txt in enumerate(marker_text):
            axes[1].annotate(txt, (marker_x[i],marker_y[i]),
                             xytext=(2, 3), textcoords='offset points')

    # legend
    for i, ax in enumerate(axes):
        p_offset = bottom_margin + (n_panels-i-1)*(h+gap)
        ax.legend(loc='center left', title=titles[i],
                  bbox_to_anchor=(legend_offset/fig_width,
                                 (p_offset + 0.5*h)/fig_height),
                  bbox_transform=fig.transFigure,
                  labelspacing=0.1, fontsize="small",
                  borderpad=0.0, borderaxespad=0.0)

    # y axis label
    fig.text(0, (bottom_margin + 0.5*panels_total_height)/fig_height,
             'concurrence after $T=\SI{200}{ns}$',
             rotation='vertical', ha='left', va='center')

    # output
    fig.savefig(outfile, format=os.path.splitext(outfile)[1][1:])

def read_data(data_folder):
    """
    Return dictionary

    run_folder => {
        'E' => array of amplitudes
        'concurrence' => array of concurrences for each value in E,
        'pop_loss' => array of population loss from log subsp (E)
        'nq_vals' => array of number of qubit levels included in prop (E),
        'nc_vals' => array of number of qubit levels included in prop (E)
    }
    """
    result = {}
    subdirs = [os.path.join(data_folder,o) for o in os.listdir(data_folder)
               if os.path.isdir(os.path.join(data_folder,o))]
    for folder in subdirs:
        foldername = os.path.split(folder)[-1]
        datfile = glob(os.path.join(folder, 'entanglement*.dat'))[0]
        try:
            E, concurrence, pop_loss = np.genfromtxt(datfile, usecols=(0,1,2),
                                                    unpack=True)
            nq_vals, nc_vals = np.genfromtxt(datfile, usecols=(3,4),
                                            dtype=np.int, unpack=True)

            result[foldername] = {
                'E':  E,
                'concurrence':  concurrence,
                'pop_loss':  pop_loss,
                'nq_vals':  nq_vals,
                'nc_vals':  nc_vals
            }
        except ValueError:
            # skip if the data file is empty (run has not finished yet)
            pass
    return result

def main(argv=None):
    if argv is None:
        argv = sys.argv
    basename = os.path.splitext(__file__)[0]
    outfile = basename + '.pdf'
    data_folder = basename
    create_figure(outfile, read_data(data_folder))

if __name__ == "__main__":
    sys.exit(main())


import logging
import os
import pathlib

import appdirs
from netCDF4 import Dataset
import numpy as np
from scipy.interpolate import griddata
from scipy.interpolate.fitpack2 import RectBivariateSpline

from pyschism.forcing.bctides.base import TidalDataProvider

logger = logging.getLogger(__name__)

TPXO_ELEVATION = 'h_tpxo9.v1.nc'
TPXO_VELOCITY = 'u_tpxo9.v1.nc'

def raise_missing_file(fpath, fname):
    raise FileNotFoundError('\n'.join([
        f'No TPXO file found at "{fpath}".',
        'New users will need to register and request a copy of '
        f'the TPXO9 NetCDF file (specifically `{fname}`) '
        'from the authors at https://www.tpxo.net.',
        'Once you obtain `h_tpxo9.v1.nc`, you can follow one of the '
        'following options: ',
        f'1) copy or symlink the file to "{fpath}"',
        f'2) set the environment variable `{fname}` to point'
        ' to the file',
    ]))

class TPXO(TidalDataProvider):

    def __init__(self, h_file=None, u_file=None):
        self._h_file = h_file
        self._u_file = u_file

    def get_elevation(self, constituent, vertices):
        logger.info('Querying TPXO for elevation constituent '
                    f'{constituent}.')
        amp = self._get_interpolation(
            'elevation', 'ha', constituent, vertices)
        phase = self._get_interpolation(
            'elevation', 'hp', constituent, vertices)
        return amp, phase

    def get_velocity(self, constituent, vertices):
        logger.info('Querying TPXO for velocity constituent '
                    f'{constituent}.')
        uamp = self._get_interpolation(
            'velocity', 'ua', constituent, vertices) / 100.
        uphase = self._get_interpolation(
            'velocity', 'up', constituent, vertices)
        vamp = self._get_interpolation(
            'velocity', 'va', constituent, vertices) / 100.
        vphase = self._get_interpolation(
            'velocity', 'vp', constituent, vertices)
        return uamp, uphase, vamp, vphase

    @property
    def constituents(self):
        return ['M2', 'S2', 'N2', 'K2', 'K1', 'O1', 'P1',
                'Q1', 'Mm', 'Mf', 'M4', 'MN4', 'MS4', '2N2', 'S1']
        if not hasattr(self, '_constituents'):
            self._constituents = [
                c.capitalize() for c in self.h['con'][:].astype(
                    '|S1').tobytes().decode('utf-8').split()]
        return self._constituents

    @property
    def x(self) -> np.ndarray:
        return self.h['lon_z'][:, 0].data

    @property
    def y(self) -> np.ndarray:
        return self.h['lat_z'][0, :].data

    @property
    def h(self):
        if not hasattr(self, '_h'):
            if self._h_file is None:
                self._h_file = os.getenv('TPXO_ELEVATION')
                if self._h_file is None:
                    self._h_file = pathlib.Path(
                        appdirs.user_data_dir('tpxo')) / TPXO_ELEVATION
            if not self._h_file.exists():
                raise_missing_file(self._h_file, TPXO_ELEVATION)
            self._h = Dataset(self._h_file)
        return self._h

    @property
    def uv(self):
        if not hasattr(self, '_uv'):
            if self._u_file is None:
                self._u_file = os.getenv('TPXO_VELOCITY')
                if self._u_file is None:
                    self._u_file = pathlib.Path(
                        appdirs.user_data_dir('tpxo')) / TPXO_VELOCITY
            if not self._u_file.exists():
                raise_missing_file(self._u_file, TPXO_VELOCITY)
            self._uv = Dataset(self._u_file)
        return self._uv

    def _get_interpolation(self, phys_var, ncvar, constituent, vertices):
        lower_c = [c.lower() for c in self.constituents]
        if phys_var == 'elevation':
            ncarray = self.h
        elif phys_var == 'velocity':
            ncarray = self.uv
        zi = ncarray[ncvar][
            lower_c.index(constituent.lower()), :, :]
        xo = np.asarray(
            [x + 360. if x < 0. else x for x in vertices[:, 0]]).flatten()
        yo = vertices[:, 1].flatten()
        xi, yi = np.meshgrid(self.x, self.y, indexing='ij')
        xi = xi.flatten()
        yi = yi.flatten()
        zi = zi.flatten()
        dx = np.mean(np.diff(self.x))
        dy = np.mean(np.diff(self.y))
        # buffer the bbox by 2 difference units
        mask1 = np.logical_and(
            np.logical_and(
                xi >= np.min(xo) - 2 * dx,
                xi <= np.max(xo) + 2 * dx
            ),
            np.logical_and(
                yi >= np.min(yo) - 2 * dy,
                yi <= np.max(yo) + 2 * dy
            )
        )
        # remove junk values from input array
        mask2 = np.ma.masked_where(zi != 0., zi)
        iidx = np.where(np.logical_and(mask1, mask2))
        values = griddata(
            (xi[iidx], yi[iidx]),
            zi[iidx],
            (xo, yo),
            method='linear',
            fill_value=np.nan,
        )
        nan_idxs = np.where(np.isnan(values))
        values[nan_idxs] = griddata(
            (xi[iidx], yi[iidx]),
            zi[iidx],
            (xo[nan_idxs], yo[nan_idxs]),
            method='nearest',
        )
        return values


#!/usr/bin/env python3
"""
Multi-compartmental OLM cell example

File: olm-example.py

Copyright 2021 NeuroML contributors
Authors: Padraig Gleeson, Ankur Sinha
"""

from neuroml import (NeuroMLDocument, IncludeType, Population, PulseGenerator, ExplicitInput, Network, SegmentGroup, Member, Property, Include, Instance, Location)
from CellBuilder import (create_cell, add_segment, add_channel_density, set_init_memb_potential, set_resistivity, set_specific_capacitance, get_seg_group_by_id)
from pyneuroml import pynml
from pyneuroml.lems import LEMSSimulation
import numpy as np

def main():
    """Main function

    Include the NeuroML model into a LEMS simulation file, run it, plot some
    data.
    """
    # Simulation bits
    sim_id = "olm_example_sim"
    simulation = LEMSSimulation(sim_id=sim_id, duration=600, dt=0.01, simulation_seed=123)
    # Include the NeuroML model file
    simulation.include_neuroml2_file(create_olm_network())
    # Assign target for the simulation
    simulation.assign_simulation_target("single_olm_cell_network")

    # Recording information from the simulation
    simulation.create_output_file(id="output0", file_name=sim_id + ".dat")
    simulation.add_column_to_output_file("output0", column_id="pop0_0_v", quantity="pop0[0]/v")
    simulation.add_column_to_output_file("output0",
                                         column_id="pop0_0_v_Seg0_soma_0",
                                         quantity="pop0/0/olm/0/v")
    simulation.add_column_to_output_file("output0",
                                         column_id="pop0_0_v_Seg1_soma_0",
                                         quantity="pop0/0/olm/1/v")
    simulation.add_column_to_output_file("output0",
                                         column_id="pop0_0_v_Seg0_axon_0",
                                         quantity="pop0/0/olm/2/v")
    simulation.add_column_to_output_file("output0",
                                         column_id="pop0_0_v_Seg1_axon_0",
                                         quantity="pop0/0/olm/3/v")
    simulation.add_column_to_output_file("output0",
                                         column_id="pop0_0_v_Seg0_dend_0",
                                         quantity="pop0/0/olm/4/v")
    simulation.add_column_to_output_file("output0",
                                         column_id="pop0_0_v_Seg1_dend_0",
                                         quantity="pop0/0/olm/6/v")
    simulation.add_column_to_output_file("output0",
                                         column_id="pop0_0_v_Seg0_dend_1",
                                         quantity="pop0/0/olm/5/v")
    simulation.add_column_to_output_file("output0",
                                         column_id="pop0_0_v_Seg1_dend_1",
                                         quantity="pop0/0/olm/7/v")
    # Save LEMS simulation to file
    sim_file = simulation.save_to_file()

    # Run the simulation using the NEURON simulator
    pynml.run_lems_with_jneuroml_neuron(sim_file, max_memory="2G", nogui=True,
                                        plot=False, skip_run=False)
    # Plot the data
    plot_data(sim_id)

def plot_data(sim_id):
    """Plot the sim data.

    Load the data from the file and plot the graph for the membrane potential
    using the pynml generate_plot utility function.

    :sim_id: ID of simulaton

    """
    data_array = np.loadtxt(sim_id + ".dat")
    pynml.generate_plot([data_array[:, 0]], [data_array[:, 1]], "Membrane potential (soma seg 0)", show_plot_already=False, save_figure_to=sim_id + "_seg0_soma0-v.png", xaxis="time (s)", yaxis="membrane potential (V)")
    pynml.generate_plot([data_array[:, 0]], [data_array[:, 2]], "Membrane potential (soma seg 1)", show_plot_already=False, save_figure_to=sim_id + "_seg1_soma0-v.png", xaxis="time (s)", yaxis="membrane potential (V)")
    pynml.generate_plot([data_array[:, 0]], [data_array[:, 3]], "Membrane potential (axon seg 0)", show_plot_already=False, save_figure_to=sim_id + "_seg0_axon0-v.png", xaxis="time (s)", yaxis="membrane potential (V)")
    pynml.generate_plot([data_array[:, 0]], [data_array[:, 4]], "Membrane potential (axon seg 1)", show_plot_already=False, save_figure_to=sim_id + "_seg1_axon0-v.png", xaxis="time (s)", yaxis="membrane potential (V)")

def create_olm_network():
    """Create the network

    :returns: name of network nml file
    """
    net_doc = NeuroMLDocument(id="network",
                              notes="OLM cell network")
    net_doc_fn = "olm_example_net.nml"
    net_doc.includes.append(IncludeType(href=create_olm_cell()))
    # Create a population: convenient to create many cells of the same type
    pop = Population(id="pop0", notes="A population for our cell",
                     component="olm", size=1, type="populationList")
    pop.instances.append(Instance(id=1, location=Location(0., 0., 0.)))
    # Input
    pulsegen = PulseGenerator(id="pg_olm", notes="Simple pulse generator", delay="100ms", duration="100ms", amplitude="0.08nA")

    exp_input = ExplicitInput(target="pop0[0]", input="pg_olm")

    net = Network(id="single_olm_cell_network", note="A network with a single population")
    net_doc.pulse_generators.append(pulsegen)
    net.explicit_inputs.append(exp_input)
    net.populations.append(pop)
    net_doc.networks.append(net)

    pynml.write_neuroml2_file(nml2_doc=net_doc, nml2_file_name=net_doc_fn, validate=True)
    return net_doc_fn

def create_olm_cell():
    """Create the complete cell.

    :returns: cell object
    """
    nml_cell_doc = NeuroMLDocument(id="oml_cell")
    cell = create_cell("olm")
    nml_cell_file = cell.id + ".cell.nml"

    # Add two soma segments
    diam = 10.0
    soma_0 = add_segment(cell,
                         prox=[0.0, 0.0, 0.0, diam],
                         dist=[0.0, 10., 0.0, diam],
                         name="Seg0_soma_0",
                         group="soma_0")

    soma_1 = add_segment(cell,
                         prox=None,
                         dist=[0.0, 10. + 10., 0.0, diam],
                         name="Seg1_soma_0",
                         parent=soma_0,
                         group="soma_0")

    # Add axon segments
    diam = 1.5
    axon_0 = add_segment(cell,
                         prox=[0.0, 0.0, 0.0, diam],
                         dist=[0.0, -75, 0.0, diam],
                         name="Seg0_axon_0",
                         parent=soma_0,
                         fraction_along=0.0,
                         group="axon_0")
    axon_1 = add_segment(cell,
                         prox=None,
                         dist=[0.0, -150, 0.0, diam],
                         name="Seg1_axon_0",
                         parent=axon_0,
                         group="axon_0")

    # Add 2 dendrite segments

    diam = 3.0
    dend_0_0 = add_segment(cell,
                           prox=[0.0, 20, 0.0, diam],
                           dist=[100, 120, 0.0, diam],
                           name="Seg0_dend_0",
                           parent=soma_1,
                           fraction_along=1,
                           group="dend_0")

    dend_1_0 = add_segment(cell,
                           prox=None,
                           dist=[177, 197, 0.0, diam],
                           name="Seg1_dend_0",
                           parent=dend_0_0,
                           fraction_along=1,
                           group="dend_0")

    dend_0_1 = add_segment(cell,
                           prox=[0.0, 20, 0.0, diam],
                           dist=[-100, 120, 0.0, diam],
                           name="Seg0_dend_1",
                           parent=soma_1,
                           fraction_along=1,
                           group="dend_1")
    dend_1_1 = add_segment(cell,
                           prox=None,
                           dist=[-177, 197, 0.0, diam],
                           name="Seg1_dend_1",
                           parent=dend_0_1,
                           fraction_along=1,
                           group="dend_1")

    # XXX: For segment groups to be correctly mapped to sections in NEURON,
    # they must include the correct neurolex ID
    for section_name in ["soma_0", "axon_0", "dend_0", "dend_1"]:
        section_group = get_seg_group_by_id(section_name, cell)
        section_group.neuro_lex_id = 'sao864921383'

    den_seg_group = get_seg_group_by_id("dendrite_group", cell)
    den_seg_group.includes.append(Include(segment_groups="dend_0"))
    den_seg_group.includes.append(Include(segment_groups="dend_1"))
    den_seg_group.properties.append(Property(tag="color", value="0.8 0 0"))

    ax_seg_group = get_seg_group_by_id("axon_group", cell)
    ax_seg_group.includes.append(Include(segment_groups="axon_0"))
    ax_seg_group.properties.append(Property(tag="color", value="0 0.8 0"))

    soma_seg_group = get_seg_group_by_id("soma_group", cell)
    soma_seg_group.includes.append(Include(segment_groups="soma_0"))

    soma_seg_group.properties.append(Property(tag="color", value="0 0 0.8"))

    # Other cell properties
    set_init_memb_potential(cell, "-67mV")
    set_resistivity(cell, "0.15 kohm_cm")
    set_specific_capacitance(cell, "1.3 uF_per_cm2")

    # channels
    # leak
    add_channel_density(cell, nml_cell_doc,
                        cd_id="leak_all",
                        cond_density="0.01 mS_per_cm2",
                        ion_channel="leak_chan",
                        ion_chan_def_file="olm-example/leak_chan.channel.nml",
                        erev="-67mV",
                        ion="non_specific")
    # HCNolm_soma
    add_channel_density(cell, nml_cell_doc,
                        cd_id="HCNolm_soma",
                        cond_density="0.5 mS_per_cm2",
                        ion_channel="HCNolm",
                        ion_chan_def_file="olm-example/HCNolm.channel.nml",
                        erev="-32.9mV",
                        ion="h",
                        group="soma_group")
    # Kdrfast_soma
    add_channel_density(cell, nml_cell_doc,
                        cd_id="Kdrfast_soma",
                        cond_density="73.37 mS_per_cm2",
                        ion_channel="Kdrfast",
                        ion_chan_def_file="olm-example/Kdrfast.channel.nml",
                        erev="-77mV",
                        ion="k",
                        group="soma_group")
    # Kdrfast_dendrite
    add_channel_density(cell, nml_cell_doc,
                        cd_id="Kdrfast_dendrite",
                        cond_density="105.8 mS_per_cm2",
                        ion_channel="Kdrfast",
                        ion_chan_def_file="olm-example/Kdrfast.channel.nml",
                        erev="-77mV",
                        ion="k",
                        group="dendrite_group")
    # Kdrfast_axon
    add_channel_density(cell, nml_cell_doc,
                        cd_id="Kdrfast_axon",
                        cond_density="117.392 mS_per_cm2",
                        ion_channel="Kdrfast",
                        ion_chan_def_file="olm-example/Kdrfast.channel.nml",
                        erev="-77mV",
                        ion="k",
                        group="axon_group")
    # KvAolm_soma
    add_channel_density(cell, nml_cell_doc,
                        cd_id="KvAolm_soma",
                        cond_density="4.95 mS_per_cm2",
                        ion_channel="KvAolm",
                        ion_chan_def_file="olm-example/KvAolm.channel.nml",
                        erev="-77mV",
                        ion="k",
                        group="soma_group")
    # KvAolm_dendrite
    add_channel_density(cell, nml_cell_doc,
                        cd_id="KvAolm_dendrite",
                        cond_density="2.8 mS_per_cm2",
                        ion_channel="KvAolm",
                        ion_chan_def_file="olm-example/KvAolm.channel.nml",
                        erev="-77mV",
                        ion="k",
                        group="dendrite_group")
    # Nav_soma
    add_channel_density(cell, nml_cell_doc,
                        cd_id="Nav_soma",
                        cond_density="10.7 mS_per_cm2",
                        ion_channel="Nav",
                        ion_chan_def_file="olm-example/Nav.channel.nml",
                        erev="50mV",
                        ion="na",
                        group="soma_group")
    # Nav_dendrite
    add_channel_density(cell, nml_cell_doc,
                        cd_id="Nav_dendrite",
                        cond_density="23.4 mS_per_cm2",
                        ion_channel="Nav",
                        ion_chan_def_file="olm-example/Nav.channel.nml",
                        erev="50mV",
                        ion="na",
                        group="dendrite_group")
    # Nav_axon
    add_channel_density(cell, nml_cell_doc,
                        cd_id="Nav_axon",
                        cond_density="17.12 mS_per_cm2",
                        ion_channel="Nav",
                        ion_chan_def_file="olm-example/Nav.channel.nml",
                        erev="50mV",
                        ion="na",
                        group="axon_group")

    nml_cell_doc.cells.append(cell)
    pynml.write_neuroml2_file(nml_cell_doc, nml_cell_file, True, True)
    return nml_cell_file

if __name__ == "__main__":
    main()