# Copyright 2018-2020 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This module contains functions to read the structure of molecules, build a Hartree-Fock state,
build an active space and generate single and double excitations.
"""
# pylint: disable=too-many-locals
import os
from shutil import copyfile

from pennylane import numpy as np

# Bohr-Angstrom correlation coefficient (https://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0)
bohr_angs = 0.529177210903

def read_structure(filepath, outpath="."):
    r"""Read the structure of the polyatomic system from a file and returns
    a list with the symbols of the atoms in the molecule and a 1D array
    with their positions :math:`[x_1, y_1, z_1, x_2, y_2, z_2, \dots]` in
    atomic units (Bohr radius = 1).

    The atomic coordinates in the file must be in Angstroms.
    The `xyz <https://en.wikipedia.org/wiki/XYZ_file_format>`_ format is supported. Additionally,
    the new file ``structure.xyz``, containing the input geometry, is created in a directory with
    path given by ``outpath``.

    Args:
        filepath (str): name of the molecular structure file in the working directory
            or the absolute path to the file if it is located in a different folder
        outpath (str): path to the output directory

    Returns:
        tuple[list, array]: symbols of the atoms in the molecule and a 1D array with their
        positions in atomic units.

    **Example**

    >>> symbols, coordinates = read_structure('h2.xyz')
    >>> print(symbols, coordinates)
    ['H', 'H'] [0.    0.   -0.66140414    0.    0.    0.66140414]
    """
    file_in = filepath.strip()
    file_out = os.path.join(outpath, "structure.xyz")

    copyfile(file_in, file_out)

    symbols = []
    coordinates = []
    with open(file_out, encoding="utf-8") as f:
        for line in f.readlines()[2:]:
            symbol, x, y, z = line.split()
            symbols.append(symbol)
            coordinates.append(float(x))
            coordinates.append(float(y))
            coordinates.append(float(z))

    return symbols, np.array(coordinates) / bohr_angs

def active_space(electrons, orbitals, mult=1, active_electrons=None, active_orbitals=None):
    r"""Build the active space for a given number of active electrons and active orbitals.

    Post-Hartree-Fock (HF) electron correlation methods expand the many-body wave function
    as a linear combination of Slater determinants, commonly referred to as configurations.
    This configurations are generated by exciting electrons from the occupied to the
    unoccupied HF orbitals as sketched in the figure below. Since the number of configurations
    increases combinatorially with the number of electrons and orbitals this expansion can be
    truncated by defining an active space.

    The active space is created by classifying the HF orbitals as core, active and
    external orbitals:

    - Core orbitals are always occupied by two electrons
    - Active orbitals can be occupied by zero, one, or two electrons
    - The external orbitals are never occupied

    |

    .. figure:: ../../_static/qchem/sketch_active_space.png
        :align: center
        :width: 50%

    |

    .. note::
        The number of active *spin*-orbitals ``2*active_orbitals`` determines the number of
        qubits required to perform the quantum simulations of the electronic structure
        of the many-electron system.

    Args:
        electrons (int): total number of electrons
        orbitals (int): total number of orbitals
        mult (int): Spin multiplicity :math:`\mathrm{mult}=N_\mathrm{unpaired} + 1` for
            :math:`N_\mathrm{unpaired}` unpaired electrons occupying the HF orbitals.
            Possible values for ``mult`` are :math:`1, 2, 3, \ldots`. If not specified,
            a closed-shell HF state is assumed.
        active_electrons (int): Number of active electrons. If not specified, all electrons
            are treated as active.
        active_orbitals (int): Number of active orbitals. If not specified, all orbitals
            are treated as active.

    Returns:
        tuple: lists of indices for core and active orbitals

    **Example**

    >>> electrons = 4
    >>> orbitals = 4
    >>> core, active = active_space(electrons, orbitals, active_electrons=2, active_orbitals=2)
    >>> print(core) # core orbitals
    [0]
    >>> print(active) # active orbitals
    [1, 2]
    """
    # pylint: disable=too-many-branches

    if active_electrons is None:
        ncore_orbs = 0
        core = []
    else:
        if active_electrons <= 0:
            raise ValueError(
                f"The number of active electrons ({active_electrons}) " f"has to be greater than 0."
            )

        if active_electrons > electrons:
            raise ValueError(
                f"The number of active electrons ({active_electrons}) "
                f"can not be greater than the total "
                f"number of electrons ({electrons})."
            )

        if active_electrons < mult - 1:
            raise ValueError(
                f"For a reference state with multiplicity {mult}, "
                f"the number of active electrons ({active_electrons}) should be "
                f"greater than or equal to {mult - 1}."
            )

        if mult % 2 == 1:
            if active_electrons % 2 != 0:
                raise ValueError(
                    f"For a reference state with multiplicity {mult}, "
                    f"the number of active electrons ({active_electrons}) should be even."
                )
        else:
            if active_electrons % 2 != 1:
                raise ValueError(
                    f"For a reference state with multiplicity {mult}, "
                    f"the number of active electrons ({active_electrons}) should be odd."
                )

        ncore_orbs = (electrons - active_electrons) // 2
        core = list(range(ncore_orbs))

    if active_orbitals is None:
        active = list(range(ncore_orbs, orbitals))
    else:
        if active_orbitals <= 0:
            raise ValueError(
                f"The number of active orbitals ({active_orbitals}) " f"has to be greater than 0."
            )

        if ncore_orbs + active_orbitals > orbitals:
            raise ValueError(
                f"The number of core ({ncore_orbs}) + active orbitals ({active_orbitals}) cannot "
                f"be greater than the total number of orbitals ({orbitals})"
            )

        homo = (electrons + mult - 1) / 2
        if ncore_orbs + active_orbitals <= homo:
            raise ValueError(
                f"For n_active_orbitals={active_orbitals}, there are no virtual orbitals "
                f"in the active space."
            )

        active = list(range(ncore_orbs, ncore_orbs + active_orbitals))

    return core, active

def excitations(electrons, orbitals, delta_sz=0):
    r"""Generate single and double excitations from a Hartree-Fock reference state.

    Single and double excitations can be generated by acting with the operators
    :math:`\hat T_1` and :math:`\hat T_2` on the Hartree-Fock reference state:

    .. math::

        && \hat{T}_1 = \sum_{r \in \mathrm{occ} \\ p \in \mathrm{unocc}}
        \hat{c}_p^\dagger \hat{c}_r \\
        && \hat{T}_2 = \sum_{r>s \in \mathrm{occ} \\ p>q \in
        \mathrm{unocc}} \hat{c}_p^\dagger \hat{c}_q^\dagger \hat{c}_r \hat{c}_s.

    In the equations above the indices :math:`r, s` and :math:`p, q` run over the
    occupied (occ) and unoccupied (unocc) *spin* orbitals and :math:`\hat c` and
    :math:`\hat c^\dagger` are the electron annihilation and creation operators,
    respectively.

    |

    .. figure:: ../../_static/qchem/sd_excitations.png
        :align: center
        :width: 80%

    |

    Args:
        electrons (int): Number of electrons. If an active space is defined, this
            is the number of active electrons.
        orbitals (int): Number of *spin* orbitals. If an active space is defined,
            this is the number of active spin-orbitals.
        delta_sz (int): Specifies the selection rules ``sz[p] - sz[r] = delta_sz`` and
            ``sz[p] + sz[p] - sz[r] - sz[s] = delta_sz`` for the spin-projection ``sz`` of
            the orbitals involved in the single and double excitations, respectively.
            ``delta_sz`` can take the values :math:`0`, :math:`\pm 1` and :math:`\pm 2`.

    Returns:
        tuple(list, list): lists with the indices of the spin orbitals involved in the
        single and double excitations

    **Example**

    >>> electrons = 2
    >>> orbitals = 4
    >>> singles, doubles = excitations(electrons, orbitals)
    >>> print(singles)
    [[0, 2], [1, 3]]
    >>> print(doubles)
    [[0, 1, 2, 3]]
    """

    if not electrons > 0:
        raise ValueError(
            f"The number of active electrons has to be greater than 0 \n"
            f"Got n_electrons = {electrons}"
        )

    if orbitals <= electrons:
        raise ValueError(
            f"The number of active spin-orbitals ({orbitals}) "
            f"has to be greater than the number of active electrons ({electrons})."
        )

    if delta_sz not in (0, 1, -1, 2, -2):
        raise ValueError(
            f"Expected values for 'delta_sz' are 0, +/- 1 and +/- 2 but got ({delta_sz})."
        )

    # define the spin projection 'sz' of the single-particle states
    sz = np.array([0.5 if (i % 2 == 0) else -0.5 for i in range(orbitals)])

    singles = [
        [r, p]
        for r in range(electrons)
        for p in range(electrons, orbitals)
        if sz[p] - sz[r] == delta_sz
    ]

    doubles = [
        [s, r, q, p]
        for s in range(electrons - 1)
        for r in range(s + 1, electrons)
        for q in range(electrons, orbitals - 1)
        for p in range(q + 1, orbitals)
        if (sz[p] + sz[q] - sz[r] - sz[s]) == delta_sz
    ]

    return singles, doubles

def hf_state(electrons, orbitals):
    r"""Generate the occupation-number vector representing the Hartree-Fock state.

    The many-particle wave function in the Hartree-Fock (HF) approximation is a `Slater determinant
    <https://en.wikipedia.org/wiki/Slater_determinant>`_. In Fock space, a Slater determinant
    for :math:`N` electrons is represented by the occupation-number vector:

    .. math::

        \vert {\bf n} \rangle = \vert n_1, n_2, \dots, n_\mathrm{orbs} \rangle,
        n_i = \left\lbrace \begin{array}{ll} 1 & i \leq N \\ 0 & i > N \end{array} \right.,

    where :math:`n_i` indicates the occupation of the :math:`i`-th orbital.

    Args:
        electrons (int): Number of electrons. If an active space is defined, this
            is the number of active electrons.
        orbitals (int): Number of *spin* orbitals. If an active space is defined,
            this is the number of active spin-orbitals.

    Returns:
        array: NumPy array containing the vector :math:`\vert {\bf n} \rangle`

    **Example**

    >>> state = hf_state(2, 6)
    >>> print(state)
    [1 1 0 0 0 0]
    """

    if electrons <= 0:
        raise ValueError(
            f"The number of active electrons has to be larger than zero; "
            f"got 'electrons' = {electrons}"
        )

    if electrons > orbitals:
        raise ValueError(
            f"The number of active orbitals cannot be smaller than the number of active electrons;"
            f" got 'orbitals'={orbitals} < 'electrons'={electrons}"
        )

    state = np.where(np.arange(orbitals) < electrons, 1, 0)

    return np.array(state)

def excitations_to_wires(singles, doubles, wires=None):
    r"""Map the indices representing the single and double excitations
    generated with the function :func:`~.excitations` to the wires that
    the Unitary Coupled-Cluster (UCCSD) template will act on.

    Args:
        singles (list[list[int]]): list with the indices ``r``, ``p`` of the two qubits
            representing the single excitation
            :math:`\vert r, p \rangle = \hat{c}_p^\dagger \hat{c}_r \vert \mathrm{HF}\rangle`
        doubles (list[list[int]]): list with the indices ``s``, ``r``, ``q``, ``p`` of the four
            qubits representing the double excitation
            :math:`\vert s, r, q, p \rangle = \hat{c}_p^\dagger \hat{c}_q^\dagger
            \hat{c}_r \hat{c}_s \vert \mathrm{HF}\rangle`
        wires (Iterable[Any]): Wires of the quantum device. If None, will use consecutive wires.

    The indices :math:`r, s` and :math:`p, q` in these lists correspond, respectively, to the
    occupied and virtual orbitals involved in the generated single and double excitations.

    Returns:
        tuple[list[list[Any]], list[list[list[Any]]]]: lists with the sequence of wires,
        resulting from the single and double excitations, that the Unitary Coupled-Cluster
        (UCCSD) template will act on.

    **Example**

    >>> singles = [[0, 2], [1, 3]]
    >>> doubles = [[0, 1, 2, 3]]
    >>> singles_wires, doubles_wires = excitations_to_wires(singles, doubles)
    >>> print(singles_wires)
    [[0, 1, 2], [1, 2, 3]]
    >>> print(doubles_wires)
    [[[0, 1], [2, 3]]]

    >>> wires=['a0', 'b1', 'c2', 'd3']
    >>> singles_wires, doubles_wires = excitations_to_wires(singles, doubles, wires=wires)
    >>> print(singles_wires)
    [['a0', 'b1', 'c2'], ['b1', 'c2', 'd3']]
    >>> print(doubles_wires)
    [[['a0', 'b1'], ['c2', 'd3']]]
    """

    if (not singles) and (not doubles):
        raise ValueError(
            f"'singles' and 'doubles' lists can not be both empty; "
            f"got singles = {singles}, doubles = {doubles}"
        )

    expected_shape = (2,)
    for single_ in singles:
        if np.array(single_).shape != expected_shape:
            raise ValueError(
                f"Expected entries of 'singles' to be of shape (2,); got {np.array(single_).shape}"
            )

    expected_shape = (4,)
    for double_ in doubles:
        if np.array(double_).shape != expected_shape:
            raise ValueError(
                f"Expected entries of 'doubles' to be of shape (4,); got {np.array(double_).shape}"
            )

    max_idx = 0
    if singles:
        max_idx = np.max(singles)
    if doubles:
        max_idx = max(np.max(doubles), max_idx)

    if wires is None:
        wires = range(max_idx + 1)
    elif len(wires) != max_idx + 1:
        raise ValueError(f"Expected number of wires is {max_idx + 1}; got {len(wires)}")

    singles_wires = []
    for r, p in singles:
        s_wires = [wires[i] for i in range(r, p + 1)]
        singles_wires.append(s_wires)

    doubles_wires = []
    for s, r, q, p in doubles:
        d1_wires = [wires[i] for i in range(s, r + 1)]
        d2_wires = [wires[i] for i in range(q, p + 1)]
        doubles_wires.append([d1_wires, d2_wires])

    return singles_wires, doubles_wires


# -*- coding: utf-8 -*-

"""Vector Field.

This is some documentation.

An example construction.

.. todo::

    store VECTORFIELD_REPRESENTATIONS keys as str, not the BaseRepresentation
    class.

"""

__all__ = [
    "BaseVectorField",
    "CartesianVectorField",
    "CylindricalVectorField",
    "SphericalVectorField",
    "PhysicsSphericalVectorField",
]

##############################################################################
# IMPORTS

# BUILT-IN
import functools
import inspect
import operator
import typing as T

# THIRD PARTY
import astropy.coordinates as coord
import astropy.units as u
import numpy as np
from astropy.coordinates.representation import (
    REPRESENTATION_CLASSES as _REP_CLSs,
)
from astropy.coordinates.representation import (
    BaseRepresentationOrDifferential,
    _array2string,
    _make_getter,
)
from erfa import ufunc as erfa_ufunc

# PROJECT-SPECIFIC
from .coordinates import resolve_framelike
from discO.type_hints import (
    FrameLikeType,
    FrameType,
    QuantityType,
    RepresentationType,
)

##############################################################################
# PARAMETERS

_VECTORFIELD_CLASSES: T.Dict[str, object] = {}
VECTORFIELD_REPRESENTATIONS: T.Dict[coord.BaseRepresentation, object] = {}

def _invalidate_psp_cls_hash():
    global _PSP_HASH
    _PSP_HASH = None

##############################################################################
# CODE
##############################################################################

class BaseVectorField(BaseRepresentationOrDifferential):
    """Base Vector-Field.

    Parameters
    ----------
    points : |Representation|
    *args
        The components
    frame : frame-like or None (optional, keyword only)
        The frame of the vector-field. None (default), does not attach a frame.
    **kwargs
        passed along

    """

    def __init_subclass__(cls, **kwargs) -> None:
        """Set default ``attr_classes`` and component getters on a VectorField.
        class BaseVectorField(BaseRepresentationOrDifferential):
        For these, the components are those of the base representation prefixed
        by 'd_', and the class is `~astropy.units.Quantity`.

        """
        if not hasattr(cls, "base_representation"):
            raise NotImplementedError(
                "VectorField representations must have a"
                '"base_representation" class attribute.',
            )

        # If not defined explicitly, create attr_classes.
        if not hasattr(cls, "attr_classes"):
            base_attr_classes = cls.base_representation.attr_classes
            cls.attr_classes = {
                "vf_" + c: u.Quantity for c in base_attr_classes
            }

        # Now check caches!
        repr_name = cls.get_name()
        if repr_name in _VECTORFIELD_CLASSES:
            raise ValueError(
                f"VectorField class '{repr_name}' already exists.",
            )
        elif cls.base_representation in VECTORFIELD_REPRESENTATIONS:
            raise ValueError(
                "VectorField with representation "
                f"'{cls.base_representation}' already exists.",
            )

        _VECTORFIELD_CLASSES[repr_name] = cls
        _invalidate_psp_cls_hash()

        # add to representations dict
        VECTORFIELD_REPRESENTATIONS[cls.base_representation] = cls

        # If not defined explicitly, create properties for the components.
        for component in cls.attr_classes:
            if not hasattr(cls, component):
                setattr(
                    cls,
                    component,
                    property(
                        _make_getter(component),
                        doc=f"Component '{component}' of the VectorField.",
                    ),
                )

        super().__init_subclass__(**kwargs)

    # /def

    def __init__(
        self,
        points: RepresentationType,
        *args,
        frame: T.Optional[FrameLikeType] = None,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)
        self._frame = None if frame is None else resolve_framelike(frame)

        vf_q1 = getattr(self, "_" + self.components[0])
        vf_qs = [getattr(self, "_" + c) for c in self.components[1:]]

        if not all(vf_q1.unit.is_equivalent(vf_q.unit) for vf_q in vf_qs):
            raise u.UnitsError("components should have equivalent units.")

        if not isinstance(points, coord.BaseRepresentation):
            raise TypeError("points is not <BaseRepresentation>.")

        # TODO store in CoordinateFrame. If representation, use GenericFrame
        # "points" property and stuff links to the _points.data
        self._points = points.represent_as(self.base_representation)

    # /def

    @property
    def points(self):
        return self._points

    @property
    def frame(self) -> FrameType:
        return self._frame

    # /def

    #######################################################
    # Representation

    def to_cartesian(self):
        """Convert the field to 3D rectangular cartesian coordinates.

        Returns
        -------
        `CartesianVectorField`
            This object, converted

        """
        base_e = self.points.unit_vectors()
        c = functools.reduce(
            operator.add,
            (
                getattr(self, d_c) * base_e[c]
                for d_c, c in zip(self.components, self.points.components)
            ),
        )

        return CartesianVectorField(
            self.points.to_cartesian(),
            vf_x=c.x,
            vf_y=c.y,
            vf_z=c.z,
            frame=self.frame,
        )

    # /def

    @classmethod
    def from_cartesian(cls, other):
        """Convert field from 3D Cartesian coordinates to the desired class.

        Parameters
        ----------
        other : `CartesianVectorField`
            The object to convert into this vector field.

        Returns
        -------
        BaseVectorField
            A new Vector Field object that is this class' type.

        """
        points = cls.base_representation.from_cartesian(other.points)
        base_e = points.unit_vectors()

        return cls(
            points,
            *(other.dot(e) for e in base_e.values()),
            copy=False,
            frame=other.frame,
        )

    # /def

    def represent_as(self, other_class):
        """Convert coordinates to another representation.

        If the instance is of the requested class, it is returned unmodified.
        By default, conversion is done via cartesian coordinates.

        Parameters
        ----------
        other_class : `~BaseVectorField` subclass
            The type of representation to turn the coordinates into.

        """
        if other_class is self.__class__:
            return self

        # The default is to convert via cartesian coordinates.
        self_cartesian = self.to_cartesian()

        if inspect.isclass(other_class) and issubclass(
            other_class,
            BaseVectorField,
        ):
            pass
        elif inspect.isclass(other_class) and issubclass(
            other_class,
            coord.BaseRepresentation,
        ):
            # convert other_class to the corresponding VectorField
            other_class = VECTORFIELD_REPRESENTATIONS[other_class]
        elif isinstance(other_class, str):
            rep_cls = _REP_CLSs[other_class]
            # convert rep_cls to the corresponding VectorField
            other_class = VECTORFIELD_REPRESENTATIONS[rep_cls]
        else:
            raise TypeError

        return other_class.from_cartesian(self_cartesian)

    # /def

    @classmethod
    def from_field(cls, vectorfield):
        """Create a new instance of this vectorfield from another one.

        Parameters
        ----------
        vectorfield : `~BaseVectorField` instance
            The presentation that should be converted to this class.

        """
        if isinstance(vectorfield, BaseVectorField):
            cartesian = vectorfield.to_cartesian(
                # base.represent_as(vectorfield.base_representation)
            )
        else:
            raise TypeError

        return cls.from_cartesian(cartesian)

    # /def

    #######################################################
    # math

    def _scale_operation(self, op: T.Callable, *args):
        """Scale all components.

        Parameters
        ----------
        op : `~operator` callable
            Operator to apply (e.g., `~operator.mul`, `~operator.neg`, etc.
        *args
            Any arguments required for the operator (typically, what is to
            be multiplied with, divided by).

        """
        scaled_attrs = [op(getattr(self, c), *args) for c in self.components]
        scaled_points = self.points._scale_operation(op, *args)
        return self.__class__(
            scaled_points,
            *scaled_attrs,
            copy=False,
            frame=self.frame,
        )

    # /def

    def _combine_operation(self, op: T.Callable, other, reverse: bool = False):
        """Combine two vector fields.

        If ``other`` is of the same phase space position type as ``self``, the
        components will simply be combined.  If ``other`` is a representation,
        it will be used as a base for which to evaluate the phase space
        position, and the result is a new representation.

        Parameters
        ----------
        op : `~operator` callable
            Operator to apply (e.g., `~operator.add`, `~operator.sub`, etc.
        other : `BaseVectorField` instance
            The other phase space position or representation.
        reverse : bool
            Whether the operands should be reversed (e.g., as we got here via
            ``self.__rsub__`` because ``self`` is a subclass of ``other``).

        """
        # ----------
        # make sure points are the same

        diff = (
            self.points.represent_as(
                coord.CartesianRepresentation,
            )
            - other.points.represent_as(coord.CartesianRepresentation)
        )

        if not np.allclose(diff.norm().value, 0):
            raise Exception("can't combine mismatching points.")

        # ----------

        if isinstance(self, type(other)):

            first, second = (self, other) if not reverse else (other, self)
            return self.__class__(
                self.points,
                *[
                    op(getattr(first, c), getattr(second, c))
                    for c in self.components
                ],
                frame=self.frame,
            )
        else:
            try:
                self_cartesian = self.to_cartesian()
            except TypeError:
                return NotImplemented

            return other._combine_operation(op, self_cartesian, not reverse)

    # /def

    def norm(self) -> QuantityType:
        """Vector norm.

        The norm is the standard Frobenius norm, i.e., the square root of the
        sum of the squares of all components with non-angular units.

        Note that any associated differentials will be dropped during this
        operation.

        Returns
        -------
        norm : `astropy.units.Quantity`
            Vector norm, with the same shape as the representation.

        """
        return np.sqrt(
            functools.reduce(
                operator.add,
                (
                    getattr(self, component) ** 2
                    for component, cls in self.attr_classes.items()
                ),
            ),
        )

    # /def

    #######################################################
    # utils

    def unit_vectors(self) -> T.Dict[str, RepresentationType]:
        r"""Cartesian unit vectors in the direction of each component.

        Given unit vectors :math:`\hat{e}_c` and scale factors :math:`f_c`,
        a change in one component of :math:`\delta c` corresponds to a change
        in representation of :math:`\delta c \times f_c \times \hat{e}_c`.

        Returns
        -------
        unit_vectors : dict of `CartesianRepresentation`
            The keys are the component names.

        """
        return self.points.unit_vectors()

    # /def

    def scale_factors(self) -> T.Dict[str, QuantityType]:
        r"""Scale factors for each component's direction.

        Given unit vectors :math:`\hat{e}_c` and scale factors :math:`f_c`,
        a change in one component of :math:`\delta c` corresponds to a change
        in representation of :math:`\delta c \times f_c \times \hat{e}_c`.

        Returns
        -------
        scale_factors : dict of `~astropy.units.Quantity`
            The keys are the component names.

        """
        return self.points.scale_factors()

    # /def

    def __repr__(self) -> str:
        prefixstr = "    "
        # TODO combine with points
        arrstr = _array2string(
            np.lib.recfunctions.merge_arrays(
                (self.points._values, self._values),
            ),
            prefix=prefixstr,
        )

        pointsunitstr = (
            ("in " + self.points._unitstr)
            if self.points._unitstr
            else "[dimensionless]"
        )
        unitstr = (
            ("in " + self._unitstr) if self._unitstr else "[dimensionless]"
        )
        return "<{} ({}) {:s} | ({}) {:s}\n{}{}>".format(
            self.__class__.__name__,
            ", ".join(self.points.components),
            pointsunitstr,
            ", ".join(self.components),
            unitstr,
            prefixstr,
            arrstr,
        )

    # /def

    def _apply(self, method: T.Union[str, T.Callable], *args, **kwargs):
        """Create a new representation or differential with ``method`` applied
        to the component data.

        In typical usage, the method is any of the shape-changing methods for
        `~numpy.ndarray` (``reshape``, ``swapaxes``, etc.), as well as those
        picking particular elements (``__getitem__``, ``take``, etc.), which
        are all defined in `~astropy.utils.shapes.ShapedLikeNDArray`. It will be
        applied to the underlying arrays (e.g., ``x``, ``y``, and ``z`` for
        `~astropy.coordinates.CartesianRepresentation`), with the results used
        to create a new instance.

        Internally, it is also used to apply functions to the components
        (in particular, `~numpy.broadcast_to`).

        Parameters
        ----------
        method : str or callable
            If str, it is the name of a method that is applied to the internal
            ``components``. If callable, the function is applied.
        args : tuple
            Any positional arguments for ``method``.
        kwargs : dict
            Any keyword arguments for ``method``.

        """
        if callable(method):

            def apply_method(array):
                return method(array, *args, **kwargs)

        else:
            apply_method = operator.methodcaller(method, *args, **kwargs)

        new = super().__new__(self.__class__)
        new._points = self.points._apply(method, *args, **kwargs)
        for component in self.components:
            setattr(
                new,
                "_" + component,
                apply_method(getattr(self, component)),
            )

        # Copy other 'info' attr only if it has actually been defined.
        # See PR #3898 for further explanation and justification, along
        # with Quantity.__array_finalize__
        if "info" in self.__dict__:
            new.info = self.info

        return new

    # /def

# /class

# -------------------------------------------------------------------

class CartesianVectorField(BaseVectorField):
    """Cartesian Vector Field."""

    _xyz = None
    _vf_xyz = None

    base_representation = coord.CartesianRepresentation

    @property
    def x(self) -> QuantityType:
        return self.points.x

    @property
    def y(self) -> QuantityType:
        return self.points.y

    @property
    def z(self) -> QuantityType:
        return self.points.z

    def __init__(
        self,
        points: RepresentationType,
        vf_x,
        vf_y=None,
        vf_z=None,
        frame: T.Optional[FrameLikeType] = None,
        copy: bool = False,
    ) -> None:
        super().__init__(points, vf_x, vf_y, vf_z, frame=frame, copy=copy)

    # /def

    def get_xyz(self, xyz_axis: int = 0) -> QuantityType:
        """Return a vector array of the x, y, and z coordinates.

        Parameters
        ----------
        xyz_axis : int, optional
            The axis in the final array along which the x, y, z components
            should be stored (default: 0).

        Returns
        -------
        xyz : `~astropy.units.Quantity`
            With dimension 3 along ``xyz_axis``.  Note that, if possible,
            this will be a view.

        """
        return self.points.get_xyz(xyz_axis=xyz_axis)

    xyz = property(get_xyz)
    # /def

    def get_vf_xyz(self, vf_xyz_axis: int = 0):
        """Return a vector array of the vf_x, vf_y, and vf_z coordinates.

        Parameters
        ----------
        vf_xyz_axis : int, optional
            The axis in the final array along which the vf_x, vf_y, vf_z
            components should be stored (default: 0).

        Returns
        -------
        vf_xyz : `~astropy.units.Quantity`
            With dimension 3 along ``vf_xyz_axis``.  Note that, if possible,
            this will be a view.

        """
        if self._vf_xyz is not None:
            if self._vf_xyz_axis == vf_xyz_axis:
                return self._vf_xyz
            else:
                return np.moveaxis(
                    self._vf_xyz,
                    self._vf_xyz_axis,
                    vf_xyz_axis,
                )

        # Create combined array.  TO DO: keep it in _xyz for repeated use?
        # But then in-place changes have to cancel it. Likely best to
        # also update components.
        return np.stack([self._vf_x, self._vf_y, self._vf_z], axis=vf_xyz_axis)

    vf_xyz = property(get_vf_xyz)
    # /def

    def dot(self, other):
        """Dot product of two vector fields.

        Note that any associated differentials will be dropped during this
        operation.

        Parameters
        ----------
        other : `BaseVectorField` or |Representation|
            If not already cartesian, it is converted.

        Returns
        -------
        dot_product : `~astropy.units.Quantity`
            The sum of the product of the x, y, and z components of ``self``
            and ``other``.

        """
        try:
            other_c = other.to_cartesian()
        except Exception:
            raise TypeError(
                "cannot only take dot product with another "
                "vector field, not a {} instance.".format(type(other)),
            )

        if isinstance(other_c, BaseVectorField):
            other_vf_xyz = other_c.get_vf_xyz(vf_xyz_axis=-1)
        else:
            other_vf_xyz = other_c.get_xyz(xyz_axis=-1)

        # erfa pdp: p-vector inner (=scalar=dot) product.
        return erfa_ufunc.pdp(self.get_vf_xyz(vf_xyz_axis=-1), other_vf_xyz)

    # /def

# /class

# -------------------------------------------------------------------

class CylindricalVectorField(BaseVectorField):
    """Cylindrical Vector Field."""

    base_representation = coord.CylindricalRepresentation

    def __init__(
        self,
        points: RepresentationType,
        vf_rho,
        vf_phi=None,
        vf_z=None,
        frame: T.Optional[FrameLikeType] = None,
        copy: bool = False,
    ) -> None:
        super().__init__(points, vf_rho, vf_phi, vf_z, frame=frame, copy=copy)

    # /def

    @property
    def rho(self) -> QuantityType:
        return self.points.rho

    @property
    def phi(self) -> QuantityType:
        return self.points.phi

    @property
    def z(self) -> QuantityType:
        return self.points.z

    # /def

# /class

# -------------------------------------------------------------------

class SphericalVectorField(BaseVectorField):
    """Spherical Vector Field."""

    base_representation = coord.SphericalRepresentation

    def __init__(
        self,
        points: RepresentationType,
        vf_lon,
        vf_lat=None,
        vf_distance=None,
        frame: T.Optional[FrameLikeType] = None,
        copy: bool = False,
    ) -> None:
        super().__init__(
            points,
            vf_lon,
            vf_lat,
            vf_distance,
            frame=frame,
            copy=copy,
        )

    # /def

    @property
    def lon(self) -> QuantityType:
        return self.points.lon

    @property
    def lat(self) -> QuantityType:
        return self.points.lat

    @property
    def distance(self) -> QuantityType:
        return self.points.distance

    # /def

# /class

# -------------------------------------------------------------------

class PhysicsSphericalVectorField(BaseVectorField):
    """PhysicsSpherical Vector Field."""

    base_representation = coord.PhysicsSphericalRepresentation

    def __init__(
        self,
        points: RepresentationType,
        vf_phi,
        vf_theta=None,
        vf_r=None,
        frame: T.Optional[FrameLikeType] = None,
        copy: bool = False,
    ) -> None:
        super().__init__(
            points,
            vf_phi,
            vf_theta,
            vf_r,
            frame=frame,
            copy=copy,
        )

    # /def

    @property
    def phi(self) -> QuantityType:
        return self.points.phi

    @property
    def theta(self) -> QuantityType:
        return self.points.theta

    @property
    def r(self) -> QuantityType:
        return self.points.r

    # /def

# /class

##############################################################################
# END

import numpy,import astropy,from astropy


# Third-party
import astropy.coordinates as coord
import astropy.units as u
import numpy as np
from scipy.signal import argrelmax

# Project
from gala.logging import logger
from .core import PhaseSpacePosition
from .util import peak_to_peak_period
from .plot import plot_projections
from ..io import quantity_to_hdf5, quantity_from_hdf5
from ..util import atleast_2d
from ..units import dimensionless, UnitSystem, DimensionlessUnitSystem

__all__ = ['Orbit']

class Orbit(PhaseSpacePosition):
    """
    Represents an orbit: positions and velocities (conjugate momenta) as a
    function of time.

    The class can be instantiated with Astropy representation objects (e.g.,
    :class:`~astropy.coordinates.CartesianRepresentation`), Astropy
    :class:`~astropy.units.Quantity` objects, or plain Numpy arrays.

    If passing in Quantity or Numpy array instances for both position and
    velocity, they are assumed to be Cartesian. Array inputs are interpreted as
    dimensionless quantities. The input position and velocity objects can have
    an arbitrary number of (broadcastable) dimensions. For Quantity or array
    inputs, the first axes have special meaning:

        - ``axis=0`` is the coordinate dimension (e.g., x, y, z)
        - ``axis=1`` is the time dimension

    So if the input position array, ``pos``, has shape ``pos.shape = (3, 100)``,
    this would be a 3D orbit at 100 times (``pos[0]`` is ``x``, ``pos[1]``` is
    ``y``, etc.). For representing multiple orbits, the position array could
    have 3 axes, e.g., it might have shape `pos.shape = (3, 100, 8)`, where this
    is interpreted as a 3D position at 100 times for 8 different orbits. The
    same is true for velocity. The position and velocity arrays must have the
    same shape.

    If a time argument is specified, the position and velocity arrays must have
    the same number of timesteps as the length of the time object::

        len(t) == pos.shape[1]

    Parameters
    ----------
    pos : representation, quantity_like, or array_like
        Positions. If a numpy array (e.g., has no units), this will be
        stored as a dimensionless :class:`~astropy.units.Quantity`. See
        the note above about the assumed meaning of the axes of this object.
    vel : differential, quantity_like, or array_like
        Velocities. If a numpy array (e.g., has no units), this will be
        stored as a dimensionless :class:`~astropy.units.Quantity`. See
        the note above about the assumed meaning of the axes of this object.
    t : array_like, :class:`~astropy.units.Quantity` (optional)
        Array of times. If a numpy array (e.g., has no units), this will be
        stored as a dimensionless :class:`~astropy.units.Quantity`.
    hamiltonian : `~gala.potential.Hamiltonian` (optional)
        The Hamiltonian that the orbit was integrated in.

    """
    def __init__(self, pos, vel, t=None,
                 hamiltonian=None, potential=None, frame=None):

        super().__init__(pos=pos, vel=vel)

        if self.pos.ndim < 1:
            self.pos = self.pos.reshape(1)
            self.vel = self.vel.reshape(1)

        # TODO: check that Hamiltonian ndim is consistent with here

        if t is not None:
            t = np.atleast_1d(t)
            if self.pos.shape[0] != len(t):
                raise ValueError("Position and velocity must have the same "
                                 "length along axis=1 as the length of the "
                                 "time array {} vs {}"
                                 .format(len(t), self.pos.shape[0]))

            if not hasattr(t, 'unit'):
                t = t * u.one

        self.t = t

        if hamiltonian is not None:
            self.potential = hamiltonian.potential
            self.frame = hamiltonian.frame

        else:
            self.potential = potential
            self.frame = frame

    def __getitem__(self, slice_):

        if isinstance(slice_, np.ndarray) or isinstance(slice_, list):
            slice_ = (slice_,)

        try:
            slice_ = tuple(slice_)
        except TypeError:
            slice_ = (slice_,)

        kw = dict()
        if self.t is not None:
            kw['t'] = self.t[slice_[0]]

        pos = self.pos[slice_]
        vel = self.vel[slice_]

        # if one time is sliced out, return a phasespaceposition
        try:
            int_tslice = int(slice_[0])
        except TypeError:
            int_tslice = None

        if int_tslice is not None:
            return PhaseSpacePosition(pos=pos, vel=vel, frame=self.frame)

        else:
            return self.__class__(pos=pos, vel=vel,
                                  potential=self.potential,
                                  frame=self.frame, **kw)

    @property
    def hamiltonian(self):
        if self.potential is None or self.frame is None:
            return None

        try:
            return self._hamiltonian
        except AttributeError:
            from gala.potential import Hamiltonian
            self._hamiltonian = Hamiltonian(potential=self.potential,
                                            frame=self.frame)

        return self._hamiltonian

    def w(self, units=None):
        """
        This returns a single array containing the phase-space positions.

        Parameters
        ----------
        units : `~gala.units.UnitSystem` (optional)
            The unit system to represent the position and velocity in
            before combining into the full array.

        Returns
        -------
        w : `~numpy.ndarray`
            A numpy array of all positions and velocities, without units.
            Will have shape ``(2*ndim, ...)``.

        """

        if units is None:
            if self.hamiltonian is None:
                units = dimensionless
            else:
                units = self.hamiltonian.units

        return super().w(units=units)

    # ------------------------------------------------------------------------
    # Convert from Cartesian to other representations
    #
    def represent_as(self, new_pos, new_vel=None):
        """
        Represent the position and velocity of the orbit in an alternate
        coordinate system. Supports any of the Astropy coordinates
        representation classes.

        Parameters
        ----------
        new_pos : :class:`~astropy.coordinates.BaseRepresentation`
            The type of representation to generate. Must be a class (not an
            instance), or the string name of the representation class.
        new_vel : :class:`~astropy.coordinates.BaseDifferential` (optional)
            Class in which any velocities should be represented. Must be a class
            (not an instance), or the string name of the differential class. If
            None, uses the default differential for the new position class.

        Returns
        -------
        new_orbit : `gala.dynamics.Orbit`
        """
        kw = dict()
        if self.t is not None:
            kw['t'] = self.t
        o = super().represent_as(new_pos=new_pos, new_vel=new_vel)
        return self.__class__(pos=o.pos,
                              vel=o.vel,
                              hamiltonian=self.hamiltonian,
                              **kw)

    # ------------------------------------------------------------------------
    # Shape and size
    # ------------------------------------------------------------------------
    @property
    def ntimes(self):
        return self.shape[0]

    @property
    def norbits(self):
        if len(self.shape) < 2:
            return 1
        else:
            return self.shape[1]

    def reshape(self, new_shape):
        """
        Reshape the underlying position and velocity arrays.
        """
        kw = dict()
        if self.t is not None:
            kw['t'] = self.t
        return self.__class__(pos=self.pos.reshape(new_shape),
                              vel=self.vel.reshape(new_shape),
                              hamiltonian=self.hamiltonian,
                              **kw)

    # ------------------------------------------------------------------------
    # Input / output
    #
    def to_hdf5(self, f):
        """
        Serialize this object to an HDF5 file.

        Requires ``h5py``.

        Parameters
        ----------
        f : str, :class:`h5py.File`
            Either the filename or an open HDF5 file.
        """

        f = super().to_hdf5(f)

        if self.potential is not None:
            import yaml
            from ..potential.potential.io import to_dict
            f['potential'] = yaml.dump(to_dict(self.potential)).encode('utf-8')

        if self.t:
            quantity_to_hdf5(f, 'time', self.t)

        return f

    @classmethod
    def from_hdf5(cls, f):
        """
        Load an object from an HDF5 file.

        Requires ``h5py``.

        Parameters
        ----------
        f : str, :class:`h5py.File`
            Either the filename or an open HDF5 file.
        """
        # TODO: this is duplicated code from PhaseSpacePosition
        if isinstance(f, str):
            import h5py
            f = h5py.File(f, mode='r')
            close = True
        else:
            close = False

        pos = quantity_from_hdf5(f['pos'])
        vel = quantity_from_hdf5(f['vel'])

        time = None
        if 'time' in f:
            time = quantity_from_hdf5(f['time'])

        frame = None
        if 'frame' in f:
            g = f['frame']

            frame_mod = g.attrs['module']
            frame_cls = g.attrs['class']
            frame_units = [u.Unit(x.decode('utf-8')) for x in g['units']]

            if u.dimensionless_unscaled in frame_units:
                units = DimensionlessUnitSystem()
            else:
                units = UnitSystem(*frame_units)

            pars = dict()
            for k in g['parameters']:
                pars[k] = quantity_from_hdf5(g['parameters/'+k])

            exec("from {0} import {1}".format(frame_mod, frame_cls))
            frame_cls = eval(frame_cls)

            frame = frame_cls(units=units, **pars)

        potential = None
        if 'potential' in f:
            import yaml
            from ..potential.potential.io import from_dict
            _dict = yaml.load(f['potential'][()].decode('utf-8'),
                              Loader=yaml.Loader)
            potential = from_dict(_dict)

        if close:
            f.close()

        return cls(pos=pos, vel=vel, t=time,
                   frame=frame, potential=potential)

    def orbit_gen(self):
        """
        Generator for iterating over each orbit.
        """
        if self.norbits == 1:
            yield self

        else:
            for i in range(self.norbits):
                yield self[:, i]

    # ------------------------------------------------------------------------
    # Computed dynamical quantities
    #

    def potential_energy(self, potential=None):
        r"""
        The potential energy *per unit mass*:

        .. math::

            E_\Phi = \Phi(\boldsymbol{q})

        Returns
        -------
        E : :class:`~astropy.units.Quantity`
            The potential energy.
        """
        if self.hamiltonian is None and potential is None:
            raise ValueError("To compute the potential energy, a potential"
                             " object must be provided!")
        if potential is None:
            potential = self.hamiltonian.potential

        return super().potential_energy(potential)

    def energy(self, hamiltonian=None):
        r"""
        The total energy *per unit mass*:

        Parameters
        ----------
        hamiltonian : `gala.potential.Hamiltonian`, `gala.potential.PotentialBase` instance
            The Hamiltonian object to evaluate the energy. If a potential is
            passed in, this assumes a static reference frame.

        Returns
        -------
        E : :class:`~astropy.units.Quantity`
            The total energy.
        """

        if self.hamiltonian is None and hamiltonian is None:
            raise ValueError("To compute the total energy, a hamiltonian"
                             " object must be provided!")

        if hamiltonian is None:
            hamiltonian = self.hamiltonian
        else:
            from gala.potential import Hamiltonian
            hamiltonian = Hamiltonian(hamiltonian)

        return hamiltonian(self)

    def _max_helper(self, arr, approximate=False):
        """
        Helper function for computing extrema (apocenter, pericenter, z_height)
        and times of extrema.

        Parameters
        ----------
        arr : `numpy.ndarray`
        """
        assert self.norbits == 1
        assert self.t[-1] > self.t[0]  # time must increase

        _ix = argrelmax(arr.value, mode='wrap')[0]
        _ix = _ix[(_ix != 0) & (_ix != (len(arr)-1))]  # remove edges
        t = self.t.value

        approx_arr = arr[_ix]
        approx_t = t[_ix]

        if approximate:
            return approx_arr, approx_t * self.t.unit

        better_times = np.zeros(_ix.shape, dtype=float)
        better_arr = np.zeros(_ix.shape, dtype=float)
        for i, j in enumerate(_ix):
            tvals = t[j-1:j+2]
            rvals = arr[j-1:j+2].value
            coeffs = np.polynomial.polynomial.polyfit(tvals, rvals, 2)
            better_times[i] = (-coeffs[1])/(2*coeffs[2])
            better_arr[i] = ((coeffs[2] * better_times[i]**2)
                             + (coeffs[1] * better_times[i]) + coeffs[0])

        return better_arr * arr.unit, better_times * self.t.unit

    def _max_return_helper(self, vals, times, return_times, reduce):
        if return_times:
            if len(vals) == 1:
                return vals[0], times[0]
            else:
                return vals, times

        elif reduce:
            return u.Quantity(vals).reshape(self.shape[1:])

        else:
            return u.Quantity(vals)

    def pericenter(self, return_times=False, func=np.mean,
                   approximate=False):
        """
        Estimate the pericenter(s) of the orbit by identifying local minima in
        the spherical radius, fitting a parabola around these local minima and
        then solving this parabola to find the pericenter(s).

        By default, this returns the mean of all local minima (pericenters). To
        get, e.g., the minimum pericenter, pass in ``func=np.min``. To get
        all pericenters, pass in ``func=None``.

        Parameters
        ----------
        func : func (optional)
            A function to evaluate on all of the identified pericenter times.
        return_times : bool (optional)
            Also return the pericenter times.
        approximate : bool (optional)
            Compute an approximate pericenter by skipping interpolation.

        Returns
        -------
        peri : float, :class:`~numpy.ndarray`
            Either a single number or an array of pericenters.
        times : :class:`~numpy.ndarray` (optional, see ``return_times``)
            If ``return_times=True``, also returns an array of the pericenter
            times.

        """

        if return_times and func is not None:
            raise ValueError("Cannot return times if reducing pericenters "
                             "using an input function. Pass `func=None` if "
                             "you want to return all individual pericenters "
                             "and times.")

        if func is None:
            reduce = False
            func = lambda x: x  # noqa
        else:
            reduce = True

        # time must increase
        if self.t[-1] < self.t[0]:
            self = self[::-1]

        vals = []
        times = []
        for orbit in self.orbit_gen():
            v, t = orbit._max_helper(-orbit.physicsspherical.r,  # pericenter
                                     approximate=approximate)
            vals.append(func(-v))  # negative for pericenter
            times.append(t)

        return self._max_return_helper(vals, times, return_times, reduce)

    def apocenter(self, return_times=False, func=np.mean,
                  approximate=False):
        """
        Estimate the apocenter(s) of the orbit by identifying local maxima in
        the spherical radius, fitting a parabola around these local maxima and
        then solving this parabola to find the apocenter(s).

        By default, this returns the mean of all local maxima (apocenters). To
        get, e.g., the largest apocenter, pass in ``func=np.max``. To get
        all apocenters, pass in ``func=None``.

        Parameters
        ----------
        func : func (optional)
            A function to evaluate on all of the identified apocenter times.
        return_times : bool (optional)
            Also return the apocenter times.
        approximate : bool (optional)
            Compute an approximate apocenter by skipping interpolation.

        Returns
        -------
        apo : float, :class:`~numpy.ndarray`
            Either a single number or an array of apocenters.
        times : :class:`~numpy.ndarray` (optional, see ``return_times``)
            If ``return_times=True``, also returns an array of the apocenter
            times.

        """

        if return_times and func is not None:
            raise ValueError("Cannot return times if reducing apocenters "
                             "using an input function. Pass `func=None` if "
                             "you want to return all individual apocenters "
                             "and times.")

        if func is None:
            reduce = False
            func = lambda x: x  # noqa
        else:
            reduce = True

        # time must increase
        if self.t[-1] < self.t[0]:
            self = self[::-1]

        vals = []
        times = []
        for orbit in self.orbit_gen():
            v, t = orbit._max_helper(orbit.physicsspherical.r,  # apocenter
                                     approximate=approximate)
            vals.append(func(v))
            times.append(t)

        return self._max_return_helper(vals, times, return_times, reduce)

    def zmax(self, return_times=False, func=np.mean,
             approximate=False):
        """
        Estimate the maximum ``z`` height of the orbit by identifying local
        maxima in the absolute value of the ``z`` position, fitting a parabola
        around these local maxima and then solving this parabola to find the
        maximum ``z`` height.

        By default, this returns the mean of all local maxima. To get, e.g., the
        largest ``z`` excursion, pass in ``func=np.max``. To get all ``z``
        maxima, pass in ``func=None``.

        Parameters
        ----------
        func : func (optional)
            A function to evaluate on all of the identified z maximum times.
        return_times : bool (optional)
            Also return the times of maximum.
        approximate : bool (optional)
            Compute approximate values by skipping interpolation.

        Returns
        -------
        zs : float, :class:`~numpy.ndarray`
            Either a single number or an array of maximum z heights.
        times : :class:`~numpy.ndarray` (optional, see ``return_times``)
            If ``return_times=True``, also returns an array of the apocenter
            times.

        """

        if return_times and func is not None:
            raise ValueError("Cannot return times if reducing "
                             "using an input function. Pass `func=None` if "
                             "you want to return all individual values "
                             "and times.")

        if func is None:
            reduce = False
            func = lambda x: x  # noqa
        else:
            reduce = True

        # time must increase
        if self.t[-1] < self.t[0]:
            self = self[::-1]

        vals = []
        times = []
        for orbit in self.orbit_gen():
            v, t = orbit._max_helper(np.abs(orbit.cylindrical.z),
                                     approximate=approximate)
            vals.append(func(v))
            times.append(t)

        return self._max_return_helper(vals, times, return_times, reduce)

    def eccentricity(self, **kw):
        r"""
        Returns the eccentricity computed from the mean apocenter and
        mean pericenter.

        .. math::

            e = \frac{r_{\rm apo} - r_{\rm per}}{r_{\rm apo} + r_{\rm per}}

        Parameters
        ----------
        **kw
            Any keyword arguments passed to ``apocenter()`` and
            ``pericenter()``. For example, ``approximate=True``.

        Returns
        -------
        ecc : float
            The orbital eccentricity.

        """
        ra = self.apocenter(**kw)
        rp = self.pericenter(**kw)
        return (ra - rp) / (ra + rp)

    def estimate_period(self, radial=True):
        """
        Estimate the period of the orbit. By default, computes the radial
        period. If ``radial==False``, this returns period estimates for
        each dimension of the orbit.

        Parameters
        ----------
        radial : bool (optional)
            What period to estimate. If ``True``, estimates the radial
            period. If ``False``, estimates period in each dimension, e.g.,
            if the orbit is 3D, along x, y, and z.

        Returns
        -------
        T : `~astropy.units.Quantity`
            The period or periods.
        """

        if self.t is None:
            raise ValueError("To compute the period, a time array is needed. "
                             "Specify a time array when creating this object.")

        if radial:
            r = self.physicsspherical.r.value
            if self.norbits == 1:
                T = u.Quantity(peak_to_peak_period(self.t, r))
            else:
                T = u.Quantity([peak_to_peak_period(self.t, r[:, n])
                                for n in range(r.shape[1])])

        else:
            raise NotImplementedError("sorry 'bout that...")

        return T

    # ------------------------------------------------------------------------
    # Misc. useful methods
    # ------------------------------------------------------------------------
    def circulation(self):
        """
        Determine which axes the Orbit circulates around by checking
        whether there is a change of sign of the angular momentum
        about an axis. Returns a 2D array with ``ndim`` integers per orbit
        point. If a box orbit, all integers will be 0. A 1 indicates
        circulation about the corresponding axis.

        TODO: clockwise / counterclockwise?

        For example, for a single 3D orbit:

        - Box and boxlet = [0, 0, 0]
        - z-axis (short-axis) tube = [0, 0, 1]
        - x-axis (long-axis) tube = [1, 0, 0]

        Returns
        -------
        circulation : :class:`numpy.ndarray`
            An array that specifies whether there is circulation about any of
            the axes of the input orbit. For a single orbit, will return a
            1D array, but for multiple orbits, the shape will be
            ``(3, norbits)``.

        """
        L = self.angular_momentum()

        # if only 2D, add another empty axis
        if L.ndim == 2:
            single_orbit = True
            L = L[..., None]
        else:
            single_orbit = False

        ndim, ntimes, norbits = L.shape

        # initial angular momentum
        L0 = L[:, 0]

        # see if at any timestep the sign has changed
        circ = np.ones((ndim, norbits))
        for ii in range(ndim):
            cnd = (np.sign(L0[ii]) != np.sign(L[ii, 1:])) | \
                  (np.abs(L[ii, 1:]).value < 1E-13)
            ix = np.atleast_1d(np.any(cnd, axis=0))
            circ[ii, ix] = 0

        circ = circ.astype(int)
        if single_orbit:
            return circ.reshape((ndim,))
        else:
            return circ

    def align_circulation_with_z(self, circulation=None):
        """
        If the input orbit is a tube orbit, this function aligns the circulation
        axis with the z axis and returns a copy.

        Parameters
        ----------
        circulation : array_like (optional)
            Array of bits that specify the axis about which the orbit
            circulates. If not provided, will compute this using
            :meth:`~gala.dynamics.Orbit.circulation`. See that method for more
            information.

        Returns
        -------
        orb : :class:`~gala.dynamics.Orbit`
            A copy of the original orbit object with circulation aligned with
            the z axis.
        """

        if circulation is None:
            circulation = self.circulation()
        circulation = atleast_2d(circulation, insert_axis=1)

        cart = self.cartesian
        pos = cart.xyz
        vel = np.vstack((cart.v_x.value[None],
                         cart.v_y.value[None],
                         cart.v_z.value[None])) * cart.v_x.unit

        if pos.ndim < 3:
            pos = pos[..., np.newaxis]
            vel = vel[..., np.newaxis]

        if (circulation.shape[0] != self.ndim or
                circulation.shape[1] != pos.shape[2]):
            raise ValueError("Shape of 'circulation' array should match the "
                             "shape of the position/velocity (minus the time "
                             "axis).")

        new_pos = pos.copy()
        new_vel = vel.copy()
        for n in range(pos.shape[2]):
            if circulation[2, n] == 1 or np.all(circulation[:, n] == 0):
                # already circulating about z or box orbit
                continue

            if sum(circulation[:, n]) > 1:
                logger.warning("Circulation about multiple axes - are you sure "
                               "the orbit has been integrated for long enough?")

            if circulation[0, n] == 1:
                circ = 0
            elif circulation[1, n] == 1:
                circ = 1
            else:
                raise RuntimeError("Should never get here...")

            new_pos[circ, :, n] = pos[2, :, n]
            new_pos[2, :, n] = pos[circ, :, n]

            new_vel[circ, :, n] = vel[2, :, n]
            new_vel[2, :, n] = vel[circ, :, n]

        return self.__class__(pos=new_pos.reshape(cart.xyz.shape),
                              vel=new_vel.reshape(cart.xyz.shape),
                              t=self.t,
                              hamiltonian=self.hamiltonian)

    def plot(self, components=None, units=None, auto_aspect=True, **kwargs):
        """
        Plot the positions in all projections. This is a wrapper around
        `~gala.dynamics.plot_projections` for fast access and quick
        visualization. All extra keyword arguments are passed to that function
        (the docstring for this function is included here for convenience).

        Parameters
        ----------
        components : iterable (optional)
            A list of component names (strings) to plot. By default, this is the
            Cartesian positions ``['x', 'y', 'z']``. To plot Cartesian
            velocities, pass in the velocity component names
            ``['v_x', 'v_y', 'v_z']``. If the representation is different, the
            component names will be different. For example, for a Cylindrical
            representation, the components are ``['rho', 'phi', 'z']`` and
            ``['v_rho', 'pm_phi', 'v_z']``.
        units : `~astropy.units.UnitBase`, iterable (optional)
            A single unit or list of units to display the components in.
        auto_aspect : bool (optional)
            Automatically enforce an equal aspect ratio.
        relative_to : bool (optional)
            Plot the values relative to this value or values.
        autolim : bool (optional)
            Automatically set the plot limits to be something sensible.
        axes : array_like (optional)
            Array of matplotlib Axes objects.
        subplots_kwargs : dict (optional)
            Dictionary of kwargs passed to :func:`~matplotlib.pyplot.subplots`.
        labels : iterable (optional)
            List or iterable of axis labels as strings. They should correspond to
            the dimensions of the input orbit.
        plot_function : callable (optional)
            The ``matplotlib`` plot function to use. By default, this is
            :func:`~matplotlib.pyplot.scatter`, but can also be, e.g.,
            :func:`~matplotlib.pyplot.plot`.
        **kwargs
            All other keyword arguments are passed to the ``plot_function``.
            You can pass in any of the usual style kwargs like ``color=...``,
            ``marker=...``, etc.

        Returns
        -------
        fig : `~matplotlib.Figure`

        """
        from gala.tests.optional_deps import HAS_MATPLOTLIB
        if not HAS_MATPLOTLIB:
            raise ImportError('matplotlib is required for visualization.')
        import matplotlib.pyplot as plt

        if components is None:
            if self.ndim == 1:  # only a 1D orbit, so just plot time series
                components = ['t', self.pos.components[0]]
            else:
                components = self.pos.components

        x, labels = self._plot_prepare(components=components,
                                       units=units)

        kwargs.setdefault('marker', '')
        kwargs.setdefault('linestyle', '-')
        kwargs.setdefault('labels', labels)
        kwargs.setdefault('plot_function', plt.plot)

        fig = plot_projections(x, **kwargs)

        if self.pos.get_name() == 'cartesian' and \
                all([not c.startswith('d_') for c in components]) and \
                't' not in components and \
                auto_aspect:
            for ax in fig.axes:
                ax.set(aspect='equal', adjustable='datalim')

        return fig

    def plot_3d(self, components=None, units=None, auto_aspect=True,
                subplots_kwargs=None, **kwargs):
        """
        Plot the specified 3D components.

        Parameters
        ----------
        components : iterable (optional)
            A list of component names (strings) to plot. By default, this is the
            Cartesian positions ``['x', 'y', 'z']``. To plot Cartesian
            velocities, pass in the velocity component names
            ``['v_x', 'v_y', 'v_z']``. If the representation is different, the
            component names will be different. For example, for a Cylindrical
            representation, the components are ``['rho', 'phi', 'z']`` and
            ``['v_rho', 'pm_phi', 'v_z']``.
        units : `~astropy.units.UnitBase`, iterable (optional)
            A single unit or list of units to display the components in.
        auto_aspect : bool (optional)
            Automatically enforce an equal aspect ratio.
        ax : `matplotlib.axes.Axes`
            The matplotlib Axes object to draw on.
        subplots_kwargs : dict (optional)
            Dictionary of kwargs passed to :func:`~matplotlib.pyplot.subplots`.
        labels : iterable (optional)
            List or iterable of axis labels as strings. They should correspond
            to the dimensions of the input orbit.
        plot_function : str (optional)
            The ``matplotlib`` plot function to use. By default, this is 'plot'
            but can also be, e.g., 'scatter'.
        **kwargs
            All other keyword arguments are passed to the ``plot_function``.
            You can pass in any of the usual style kwargs like ``color=...``,
            ``marker=...``, etc.

        Returns
        -------
        fig : `~matplotlib.Figure`

        """
        from gala.tests.optional_deps import HAS_MATPLOTLIB
        if not HAS_MATPLOTLIB:
            raise ImportError('matplotlib is required for visualization.')
        import matplotlib.pyplot as plt
        from mpl_toolkits import mplot3d  # noqa

        if components is None:
            components = self.pos.components

        if subplots_kwargs is None:
            subplots_kwargs = dict()

        if len(components) != 3:
            raise ValueError(
                f"The number of components ({len(components)}) must be 3")

        x, labels = self._plot_prepare(components=components,
                                       units=units)

        kwargs.setdefault('marker', '')
        kwargs.setdefault('linestyle', kwargs.pop('ls', '-'))
        plot_function_name = kwargs.pop('plot_function', 'plot')

        ax = kwargs.pop('ax', None)
        subplots_kwargs.setdefault('constrained_layout', True)
        if ax is None:
            fig, ax = plt.subplots(figsize=(6, 6),
                                   subplot_kw=dict(projection='3d'),
                                   **subplots_kwargs)
        else:
            fig = ax.figure

        plot_function = getattr(ax, plot_function_name)
        if x[0].ndim > 1:
            for n in range(x[0].shape[1]):
                plot_function(*[xx[:, n] for xx in x], **kwargs)
        else:
            plot_function(*x, **kwargs)
        ax.set_xlabel(labels[0])
        ax.set_ylabel(labels[1])
        ax.set_zlabel(labels[2])

        if self.pos.get_name() == 'cartesian' and \
                all([not c.startswith('d_') for c in components]) and \
                't' not in components and \
                auto_aspect:
            for ax in fig.axes:
                ax.set(aspect='auto', adjustable='datalim')

        return fig, ax

    def animate(self,
                components=None,
                units=None,
                stride=1,
                segment_nsteps=10,
                underplot_full_orbit=True,
                marker_style=None,
                segment_style=None,
                FuncAnimation_kwargs=None,
                orbit_plot_kwargs=None,
                axes=None):
        """
        Animate an orbit or collection of orbits.

        Parameters
        ----------
        components : iterable (optional)
            A list of component names (strings) to plot. By default, this is the
            Cartesian positions ``['x', 'y', 'z']``. To plot Cartesian
            velocities, pass in the velocity component names
            ``['v_x', 'v_y', 'v_z']``. If the representation is different, the
            component names will be different. For example, for a Cylindrical
            representation, the components are ``['rho', 'phi', 'z']`` and
            ``['v_rho', 'pm_phi', 'v_z']``.
        units : `~astropy.units.UnitBase`, iterable (optional)
            A single unit or list of units to display the components in.
        stride : int (optional)
            How often to draw a new frame, in terms of orbit timesteps.
        segment_nsteps : int (optional)
            How many timesteps to draw in an orbit segment trailing
            the timestep marker. Set this to 0 or None to disable.
        underplot_full_orbit : bool (optional)
            Controls whether to under-plot the full orbit as a thin line.
        marker_style : dict (optional)
            Matplotlib style arguments passed to `matplotlib.pyplot.plot`
            that control the plot style of the timestep marker.
        segment_style : dict (optional)
            Matplotlib style arguments passed to `matplotlib.pyplot.plot`
            that control the plot style of the orbit segment.
        FuncAnimation_kwargs : dict (optional)
            Keyword arguments passed through to
            `matplotlib.animation.FuncAnimation`.
        orbit_plot_kwargs : dict (optional)
            Keyword arguments passed through to `gala.dynamics.Orbit.plot`.
        axes : `matplotlib.axes.Axes` (optional)
            Where to draw the orbit.

        Returns
        -------
        fig : `matplotlib.figure.Figure`
        anim : `matplotlib.animation.FuncAnimation`

        """
        from gala.tests.optional_deps import HAS_MATPLOTLIB
        if not HAS_MATPLOTLIB:
            raise ImportError('matplotlib is required for visualization.')
        from matplotlib.animation import FuncAnimation

        if components is None:
            if self.ndim == 1:  # only a 1D orbit, so just plot time series
                components = ['t', self.pos.components[0]]
            else:
                components = self.pos.components

        # Extract the relevant components, in the given unit system
        xs, _ = self._plot_prepare(components=components,
                                   units=units)
        xs = [atleast_2d(xx, insert_axis=1) for xx in xs]

        # Figure out which components to plot on which axes
        data_paired = []
        for i in range(len(xs)):
            for j in range(len(xs)):
                if i >= j:
                    continue  # skip diagonal, upper triangle
                data_paired.append((xs[i], xs[j]))

        if FuncAnimation_kwargs is None:
            FuncAnimation_kwargs = dict()

        if orbit_plot_kwargs is None:
            orbit_plot_kwargs = dict()
        orbit_plot_kwargs.setdefault('zorder', 1)
        orbit_plot_kwargs.setdefault('color', '#aaaaaa')
        orbit_plot_kwargs.setdefault('linewidth', '1')
        orbit_plot_kwargs.setdefault('axes', axes)

        if marker_style is None:
            marker_style = dict()
        marker_style.setdefault('marker', 'o')
        marker_style.setdefault('linestyle', marker_style.pop('ls', 'None'))
        marker_style.setdefault('markersize', marker_style.pop('ms', 4.))
        marker_style.setdefault('color', marker_style.pop('c', 'tab:red'))
        marker_style.setdefault('zorder', 100)

        if segment_style is None:
            segment_style = dict()
        segment_style.setdefault('marker', 'None')
        segment_style.setdefault('linestyle', segment_style.pop('ls', '-'))
        segment_style.setdefault('linewidth', segment_style.pop('lw', 2.))
        segment_style.setdefault('color', segment_style.pop('c', 'tab:blue'))
        segment_style.setdefault('zorder', 10)
        if segment_nsteps is None or segment_nsteps == 0:  # HACK
            segment_style['alpha'] = 0

        # Use this to get a figure with axes with the right limits
        # Note: Labels are added by .plot()
        if not underplot_full_orbit:
            orbit_plot_kwargs['alpha'] = 0
        fig = self.plot(components=components, units=units,
                        **orbit_plot_kwargs)

        # Set up all of the (data-less) markers and line segments
        markers = []
        segments = []
        for n in range(self.norbits):
            _m = []
            _s = []
            for i in range(len(data_paired)):
                _m.append(fig.axes[i].plot([], [], **marker_style)[0])
                _s.append(fig.axes[i].plot([], [], **segment_style)[0])
            markers.append(_m)
            segments.append(_s)

        def anim_func(n):
            i = max(0, n - segment_nsteps)

            for k in range(self.norbits):
                for j in range(len(data_paired)):
                    markers[k][j].set_data(data_paired[j][0][n:n+1, k],
                                           data_paired[j][1][n:n+1, k])
                    segments[k][j].set_data(data_paired[j][0][i:n+1, k],
                                            data_paired[j][1][i:n+1, k])

            return (*[m for m in markers for x in m],
                    *[s for s in segments for x in s])

        anim = FuncAnimation(fig, anim_func,
                             frames=np.arange(0, self.ntimes, stride),
                             **FuncAnimation_kwargs)

        return fig, anim

    def to_frame(self, frame, current_frame=None, **kwargs):
        """
        Transform to a different reference frame.

        Parameters
        ----------
        frame : `gala.potential.CFrameBase`
            The frame to transform to.
        current_frame : `gala.potential.CFrameBase` (optional)
            If the Orbit has no associated Hamiltonian, this specifies the
            current frame of the orbit.

        Returns
        -------
        orbit : `gala.dynamics.Orbit`
            The orbit in the new reference frame.

        """

        kw = kwargs.copy()

        # TODO: this short-circuit sux
        if current_frame is None:
            current_frame = self.frame
        if frame == current_frame and not kwargs:
            return self

        # TODO: need a better way to do this!
        from ..potential.frame.builtin import ConstantRotatingFrame
        for fr in [frame, current_frame, self.frame]:
            if isinstance(fr, ConstantRotatingFrame):
                if 't' not in kw:
                    kw['t'] = self.t

        # TODO: this needs a re-write...
        psp = super().to_frame(frame, current_frame, **kw)

        return Orbit(pos=psp.pos, vel=psp.vel, t=self.t,
                     frame=frame, potential=self.potential)

    # ------------------------------------------------------------------------
    # Compatibility with other packages
    #

    def to_galpy_orbit(self, ro=None, vo=None):
        """Convert this object to a ``galpy.Orbit`` instance.

        Parameters
        ----------
        ro : `astropy.units.Quantity` or `astropy.units.UnitBase`
            "Natural" length unit.
        vo : `astropy.units.Quantity` or `astropy.units.UnitBase`
            "Natural" velocity unit.

        Returns
        -------
        galpy_orbit : `galpy.orbit.Orbit`

        """
        from galpy.orbit import Orbit
        from galpy.util.config import __config__ as galpy_config

        if self.frame is not None:
            from ..potential import StaticFrame
            w = self.to_frame(StaticFrame(self.frame.units))
        else:
            w = self

        if ro is None:
            ro = galpy_config.getfloat('normalization', 'ro')
            ro = ro * u.kpc

        if vo is None:
            vo = galpy_config.getfloat('normalization', 'vo')
            vo = vo * u.km/u.s

        # PhaseSpacePosition or Orbit:
        cyl = w.cylindrical

        R = cyl.rho.to_value(ro).T
        phi = cyl.phi.to_value(u.rad).T
        z = cyl.z.to_value(ro).T

        vR = cyl.v_rho.to_value(vo).T
        vT = (cyl.rho * cyl.pm_phi).to_value(vo, u.dimensionless_angles()).T
        vz = cyl.v_z.to_value(vo).T

        o = Orbit(np.array([R, vR, vT, z, vz, phi]).T, ro=ro, vo=vo)
        if w.t is not None:
            o.t = w.t.to_value(ro / vo)

        return o

    @classmethod
    def from_galpy_orbit(self, galpy_orbit):
        """Create a Gala ``PhaseSpacePosition`` or ``Orbit`` instance from a
        ``galpy.Orbit`` instance.

        Parameters
        ----------
        galpy_orbit : :class:`galpy.orbit.Orbit`

        Returns
        -------
        orbit : :class:`~gala.dynamics.Orbit`

        """
        ro = galpy_orbit._ro * u.kpc
        vo = galpy_orbit._vo * u.km/u.s
        ts = galpy_orbit.t

        rep = coord.CylindricalRepresentation(
            rho=galpy_orbit.R(ts) * ro,
            phi=galpy_orbit.phi(ts) * u.rad,
            z=galpy_orbit.z(ts) * ro
        )
        with u.set_enabled_equivalencies(u.dimensionless_angles()):
            dif = coord.CylindricalDifferential(
                d_rho=galpy_orbit.vR(ts) * vo,
                d_phi=galpy_orbit.vT(ts) * vo / rep.rho,
                d_z=galpy_orbit.vz(ts) * vo
            )

        t = galpy_orbit.t * ro / vo
        return Orbit(rep, dif, t=t)


from os import path
import numpy as np
from pvpy import PowerSpectrum

try:
    filename = 'ASTMG173.csv'
    with open(path.join(__path__[0], filename)) as file:
        pass
except IOError as e:
    print("Unable to open file %s" % filename)  # Does not exist OR no read permissions
# TODO: Maybe allow direct download of ASTM1.5G spectrum if it is missing? It shoudln't be missing.
#    import os
#    if not os.path.isfile('ASTMG173.csv'):
#        try:
#            import urllib2, StringIO, tarfile
#            data_url = 'http://rredc.nrel.gov/solar/spectra/am1.5/ASTMG173/compressed/ASTMG173.csv.tar'
#            download_as_string = urllib2.urlopen(data_url).read()
#            download_as_file = StringIO.StringIO(download_as_string)
#            download_as_tarfile_object = tarfile.open(fileobj=download_as_file)
#            download_as_tarfile_object.extractfile('ASTMG173.csv')
#        except:
#            print("Unable to open file")  # Does not exist OR no read permissions
#            raise

def jsc(*args, kind: str ="linear", spectra: str ="AM1.5G", **kwargs):
    """
    Gives the absorbed photocurrent in mA/cm**2 of a normalized spectrum.
    :param args: (array like) ideally (N,2)D numpy array with spec_in[:,0] as the wavelengths in nm and spec_in[:,1] as
     the values, but can also be two lists/vectors
    :param kind: (str) interpolation kind. see scipy.interpolate.interp1d
    :param spectra:
    :return:
    """
    if len(args) > 1:
        assert args[0].size == args[1].size, \
            "arg1 is %g, arg2 is %g. The inputs should be the same sizes." % (args[0].size, args[1].size)
        args = [np.asarray(arg) if type(arg) is list else arg for arg in args]
        spec_in = np.column_stack((args[0], args[1]))
        print("Hi")
        print(np.shape(spec_in))
        assert np.shape(spec_in)[1] == 2, "The input arguments could not be converted to a (N,2) dimensional array."
    elif len(args[0].shape) > 1 and args[0].shape[1] != 1:
        spec_in = args[0].transpose()
    else:
        spec_in = args[0]
    spec_in = np.squeeze(spec_in)
    spec = PowerSpectrum.PhotocurrentSpectrum(spec_in[0, 0], spec_in[-1, 0], spectra)
    spec.weight_spectrum(spec_in, kind=kind)
    return spec.integrate() * .1


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

"""
Written by Lucas Sinclair and Paul Rougieux.

JRC Biomass Project.
Unit D1 Bioeconomy.

Typically you can use this submodule like this:

    >>> from forest_puller.conversion.bcef_by_country import country_bcef
    >>> print(country_bcef.by_country_year)
"""

# Built-in modules #
import itertools

# Internal modules #
from forest_puller.conversion.load_expansion_factor import bcef_coefs
from forest_puller.common                           import country_codes
from forest_puller                                  import cache_dir

# First party modules #
from plumbing.cache import property_cached, property_pickled_at

# Third party modules #
import numpy, pandas

###############################################################################
class CountryBCEF:
    """
    This class uses the stock of merchantable biomass in each country
    to choose 2 factors:

    * The biomass conversion and expansion factor BCEF.
    * The root to shoot ratio R.

    These factors come from table 4.5 and 4.4 respectively in the following
    IPCC guideline document:

    https://www.ipcc-nggip.iges.or.jp/public/2006gl/pdf/4_Volume4/V4_04_Ch4_Forest_Land.pd

    We first use the BCEF_R to expand the merchantable growing stock volume to
    above-ground biomass stock. The above ground biomass stock is then used as
    a threshold to choose the root to shoot ratio.

    Interesting intermediary tables for further analysis:

    * `all_stock_merch` contains the stock in merchantable volume per ha
      and per leaf type.
    * `all_stock_abg_biomass` contains the stock in above ground biomass weight
       expressed in tons per ha and per leaf type.
    """

    min_year = 1990
    max_year = 2020

    @property_cached
    def country_climates(self):
        """
        This dataframe looks like this:

               country  climatic_zone  climatic_coef
            0       AT         boreal            0.0
            1       AT      temperate            1.0
            2       AT  mediterranean            0.0
            3       BE         boreal            0.0
            4       BE      temperate            1.0
        """
        # Load #
        df = country_codes
        # Keep only some columns #
        columns = ['iso2_code', 'boreal', 'temperate', 'mediterranean']
        df = df[columns].copy()
        # Rename column #
        df = df.rename(columns={'iso2_code': 'country'})
        # Unpivot #
        df = df.melt(id_vars    = ['country'],
                     var_name   = 'climatic_zone',
                     value_name = 'climatic_coef')
        # Sort #
        df = df.sort_values('country')
        # Reset index #
        df = df.reset_index(drop=True)
        # Return #
        return df

    @property
    def all_stock_merch(self):
        """
        This data frame looks like this:

            country  year forest_type   area stock_per_ha
        0        AT  1990         con    ...          ...
        1        AT  1990       broad    ...          ...
        2        AT  1990       mixed    ...          ...
        3        AT  2000         con    ...          ...
        4        AT  2000       broad    ...          ...
        ..      ...   ...         ...    ...          ...

        All columns are:

            ['country', 'year', 'forest_type', 'area', 'stock_per_ha']

        """
        # Import #
        import forest_puller.soef.concat
        # Load #
        area  = forest_puller.soef.concat.tables['area_by_type'].copy()
        stock = forest_puller.soef.concat.tables['stock_by_type'].copy()
        # Remove null areas and make them NaNs #
        selector = area.area == 0
        area.loc[selector, 'area'] = numpy.NaN
        # Add the area to make one big dataframe #
        df = stock.left_join(area, on=['country', 'year', 'category'])
        # Rename category to forest_type #
        df = df.rename(columns={'category': 'forest_type'})
        # Compute stock by hectare #
        df['stock_per_ha'] = df['stock'] / df['area']
        # Drop lines with NaN #
        df = df.query("stock_per_ha == stock_per_ha")
        # Now we don't need the stock column anymore #
        df = df.drop(columns=['stock'])
        # Return #
        return df

    @property
    def all_stock_merch_by_climate(self):
        """
        This dataframe looks like this:

               country  year forest_type  area   stock_per_ha  climatic_zone  climatic_coef
           0        AT  1990         con   ...            ...         boreal            0.0
           1        AT  1990         con   ...            ...      temperate            1.0
           2        AT  1990         con   ...            ...  mediterranean            0.0
           3        AT  1990       broad   ...            ...         boreal            0.0
           4        AT  1990       broad   ...            ...      temperate            1.0

        Warning: the stock per area values are duplicated for each country and forest type

        All columns are:

            ['country', 'year', 'forest_type', 'area', 'climatic_zone',
             'climatic_coef', 'stock_per_ha']
        """
        # Load #
        df = self.all_stock_merch.copy()
        # Drop mixed forest #
        df = df.query("forest_type != 'mixed'")
        # Add country info #
        df = df.left_join(self.country_climates, on="country")
        # Return #
        return df

    def get_one_bcef(self, row, kind):
        """Function to be applied to each row of the previous dataframe."""
        # If we get a NaN we return a NaN #
        if row['stock_per_ha'] != row['stock_per_ha']: return numpy.nan
        # Load #
        df = bcef_coefs
        # Select corresponding climatic zone #
        df = df.query(f"climatic_zone == '{row['climatic_zone']}'")
        # Select corresponding fores type#
        df = df.query(f"forest_type == '{row['forest_type']}'")
        # Select corresponding bounds on stock per hectare #
        df = df.query(f"lower < {row['stock_per_ha']} <= upper")
        # Make sure we have note more than one line #
        assert len(df) <= 1
        # Extract single float #
        result = df['bcef' + kind].iloc[0]
        # Return #
        return result

    @property
    def with_bcef_coefs(self):
        """
        This dataframe is the same as above except we have added three
        columns. All columns are:

            ['country', 'year', 'forest_type', 'area', 'climatic_zone',
             'climatic_coef', 'bcefi', 'bcefr', 'bcefs']
        """
        # Load #
        df = self.all_stock_merch_by_climate.copy()
        # Add three columns #
        df['bcefi'] = df.apply(lambda row: self.get_one_bcef(row, 'i'), axis=1)
        df['bcefr'] = df.apply(lambda row: self.get_one_bcef(row, 'r'), axis=1)
        df['bcefs'] = df.apply(lambda row: self.get_one_bcef(row, 's'), axis=1)
        # Now we don't need the stock_per_ha column anymore #
        df = df.drop(columns=['stock_per_ha'])
        # Return #
        return df

    @property_pickled_at('cache_path')
    def by_country_year(self):
        """
        This dataframe has three coefficients 'bcefi', 'bcefr', 'bcefs'
        for every country and for every SOEF year (except 2015).

        The dataframe looks like this:

                             bcefi     bcefr     bcefs
            country year
            AT      1990  0.600000  0.843959  0.764714
                    2000  0.574885  0.795115  0.720929
                    2005  0.573515  0.796485  0.722071
                    2010  0.572161  0.797839  0.723199
        """
        # Load #
        df = self.with_bcef_coefs.copy()
        # Multiply by the climatic situation #
        df['bcefi'] *= df['climatic_coef']
        df['bcefr'] *= df['climatic_coef']
        df['bcefs'] *= df['climatic_coef']
        # Now we don't need that column anymore #
        df = df.drop(columns=['climatic_coef'])
        # Group and sum each BCEF while keeping area #
        groups = df.groupby(['country', 'year', 'forest_type'])
        df     = groups.agg({'bcefi': 'sum',
                             'bcefr': 'sum',
                             'bcefs': 'sum',
                             'area':  'first'})
        # Get the ratio of conifers against broadleaved #
        groups           = df.groupby(['country', 'year'])
        df['area_total'] = groups['area'].transform('sum')
        df['tree_coef']  = df['area'] / df['area_total']
        # Multiply by the ratio of the given leaf type #
        df['bcefi'] *= df['tree_coef']
        df['bcefr'] *= df['tree_coef']
        df['bcefs'] *= df['tree_coef']
        # Group and sum each BCEF #
        groups = df.groupby(['country', 'year'])
        df     = groups.agg({'bcefi': 'sum',
                             'bcefr': 'sum',
                             'bcefs': 'sum'})
        df = df.reset_index()
        # Return #
        return df

    #------------------------------- Interpolation ---------------------------#
    @property
    def by_country_year_intrpld(self):
        """
        Same as above but interpolate the coefficients to get more years.
        """
        # Create a small data frame with all country and years #
        countries   = self.by_country_year['country'].drop_duplicates()
        years       = range(self.min_year, self.max_year)
        expand_grid = list(itertools.product(countries, years))
        df          = pandas.DataFrame(expand_grid, columns=('country', 'year'))
        # Join the BCEF data #
        df = df.left_join(self.by_country_year, on=['country','year'])
        # Interpolate #
        country_groups = df.groupby('country')
        df['bcefi'] = country_groups['bcefi'].transform(pandas.DataFrame.interpolate,
                                                        limit_direction='both')
        df['bcefr'] = country_groups['bcefr'].transform(pandas.DataFrame.interpolate,
                                                        limit_direction='both')
        df['bcefs'] = country_groups['bcefs'].transform(pandas.DataFrame.interpolate,
                                                        limit_direction='both')
        # Return #
        return df

    #---------------------------- Special properties -------------------------#
    @property_cached
    def all_stock_abg_biomass(self):
        """
        This data frame contains the above ground biomass stock per hectare
        expressed in tons of dry biomass. The method converts merchantable
        biomass volume (m3 of trunk) to above ground biomass weight (tons of
        dry biomass of trunk plus branches).

        The table looks like this:

            country  year forest_type  area  stock_per_ha
        1        AT  1990         con   ...           ...
        1        AT  1990       broad   ...           ...
        2        AT  1990       mixed   ...           ...
        3        AT  2000         con   ...           ...
        """
        # Load #
        stock_merch          = self.all_stock_merch
        bcef_by_country_year = self.by_country_year
        # Join the biomass conversion factors (bcef) to the stock data
        index = ['country', 'year']
        df    = stock_merch.left_join(bcef_by_country_year, on=index)
        # Compute the above ground biomass stock
        df['stock_per_ha'] *= df['bcefs']
        # Drop the coefficients columns #
        df = df.drop(columns=['bcefi', 'bcefr', 'bcefs'])
        # Return #
        return df

    # --------------------------------- Cache --------------------------------- #
    @property
    def cache_path(self):
        """Specify where on the file system we will pickle the dataframe."""
        path = cache_dir + 'conversion/bcef.pickle'
        return path

###############################################################################
# Create a singleton #
country_bcef = CountryBCEF()