# Licensed under a 3-clause BSD style license - see LICENSE.rst
import numpy as np
import astropy.units as u
from astropy.io import fits
from astropy.table import Table
from gammapy.maps import MapAxis
from gammapy.maps.utils import edges_from_lo_hi
from gammapy.utils.nddata import NDDataArray
from gammapy.utils.scripts import make_path

__all__ = ["EffectiveAreaTable", "EffectiveAreaTable2D"]

class EffectiveAreaTable:
    """Effective area table.

    TODO: Document

    Parameters
    ----------
    energy_lo, energy_hi : `~astropy.units.Quantity`
        Energy axis bin edges
    data : `~astropy.units.Quantity`
        Effective area

    Examples
    --------
    Plot parametrized effective area for HESS, HESS2 and CTA.

    .. plot::
        :include-source:

        import numpy as np
        import matplotlib.pyplot as plt
        import astropy.units as u
        from gammapy.irf import EffectiveAreaTable

        energy = np.logspace(-3, 3, 100) * u.TeV

        for instrument in ['HESS', 'HESS2', 'CTA']:
            aeff = EffectiveAreaTable.from_parametrization(energy, instrument)
            ax = aeff.plot(label=instrument)

        ax.set_yscale('log')
        ax.set_xlim([1e-3, 1e3])
        ax.set_ylim([1e3, 1e12])
        plt.legend(loc='best')
        plt.show()

    Find energy where the effective area is at 10% of its maximum value

    >>> import numpy as np
    >>> import astropy.units as u
    >>> from gammapy.irf import EffectiveAreaTable
    >>> energy = np.logspace(-1, 2) * u.TeV
    >>> aeff_max = aeff.max_area
    >>> print(aeff_max).to('m2')
    156909.413371 m2
    >>> energy_threshold = aeff.find_energy(0.1 * aeff_max)
    >>> print(energy_threshold)
    0.185368478744 TeV
    """

    def __init__(self, energy_lo, energy_hi, data, meta=None):

        e_edges = edges_from_lo_hi(energy_lo, energy_hi)
        energy_axis = MapAxis.from_edges(e_edges, interp="log", name="energy")

        interp_kwargs = {"extrapolate": False, "bounds_error": False}
        self.data = NDDataArray(
            axes=[energy_axis], data=data, interp_kwargs=interp_kwargs
        )
        self.meta = meta or {}

    @property
    def energy(self):
        return self.data.axis("energy")

    def plot(self, ax=None, energy=None, show_energy=None, **kwargs):
        """Plot effective area.

        Parameters
        ----------
        ax : `~matplotlib.axes.Axes`, optional
            Axis
        energy : `~astropy.units.Quantity`
            Energy nodes
        show_energy : `~astropy.units.Quantity`, optional
            Show energy, e.g. threshold, as vertical line

        Returns
        -------
        ax : `~matplotlib.axes.Axes`
            Axis
        """
        import matplotlib.pyplot as plt

        ax = plt.gca() if ax is None else ax

        kwargs.setdefault("lw", 2)

        if energy is None:
            energy = self.energy.center

        eff_area = self.data.evaluate(energy=energy)

        xerr = (
            (energy - self.energy.edges[:-1]).value,
            (self.energy.edges[1:] - energy).value,
        )

        ax.errorbar(energy.value, eff_area.value, xerr=xerr, **kwargs)
        if show_energy is not None:
            ener_val = u.Quantity(show_energy).to_value(self.energy.unit)
            ax.vlines(ener_val, 0, 1.1 * self.max_area.value, linestyles="dashed")
        ax.set_xscale("log")
        ax.set_xlabel(f"Energy [{self.energy.unit}]")
        ax.set_ylabel(f"Effective Area [{self.data.data.unit}]")

        return ax

    @classmethod
    def from_parametrization(cls, energy, instrument="HESS"):
        r"""Create parametrized effective area.

        Parametrizations of the effective areas of different Cherenkov
        telescopes taken from Appendix B of Abramowski et al. (2010), see
        https://ui.adsabs.harvard.edu/abs/2010MNRAS.402.1342A .

        .. math::
            A_{eff}(E) = g_1 \left(\frac{E}{\mathrm{MeV}}\right)^{-g_2}\exp{\left(-\frac{g_3}{E}\right)}

        Parameters
        ----------
        energy : `~astropy.units.Quantity`
            Energy binning, analytic function is evaluated at log centers
        instrument : {'HESS', 'HESS2', 'CTA'}
            Instrument name
        """
        energy = u.Quantity(energy)
        # Put the parameters g in a dictionary.
        # Units: g1 (cm^2), g2 (), g3 (MeV)
        # Note that whereas in the paper the parameter index is 1-based,
        # here it is 0-based
        pars = {
            "HESS": [6.85e9, 0.0891, 5e5],
            "HESS2": [2.05e9, 0.0891, 1e5],
            "CTA": [1.71e11, 0.0891, 1e5],
        }

        if instrument not in pars.keys():
            ss = f"Unknown instrument: {instrument}\n"
            ss += "Valid instruments: HESS, HESS2, CTA"
            raise ValueError(ss)

        xx = MapAxis.from_edges(energy, interp="log").center.to_value("MeV")

        g1 = pars[instrument][0]
        g2 = pars[instrument][1]
        g3 = -pars[instrument][2]

        value = g1 * xx ** (-g2) * np.exp(g3 / xx)
        data = u.Quantity(value, "cm2", copy=False)

        return cls(energy_lo=energy[:-1], energy_hi=energy[1:], data=data)

    @classmethod
    def from_constant(cls, energy, value):
        """Create constant value effective area.

        Parameters
        ----------
        energy : `~astropy.units.Quantity`
            Energy binning, analytic function is evaluated at log centers
        value : `~astropy.units.Quantity`
            Effective area
        """
        data = np.ones((len(energy) - 1)) * u.Quantity(value)
        return cls(energy_lo=energy[:-1], energy_hi=energy[1:], data=data)

    @classmethod
    def from_table(cls, table):
        """Create from `~astropy.table.Table` in ARF format.

        Data format specification: :ref:`gadf:ogip-arf`
        """
        energy_lo = table["ENERG_LO"].quantity
        energy_hi = table["ENERG_HI"].quantity
        data = table["SPECRESP"].quantity
        return cls(energy_lo=energy_lo, energy_hi=energy_hi, data=data)

    @classmethod
    def from_hdulist(cls, hdulist, hdu="SPECRESP"):
        """Create from `~astropy.io.fits.HDUList`."""
        return cls.from_table(Table.read(hdulist[hdu]))

    @classmethod
    def read(cls, filename, hdu="SPECRESP"):
        """Read from file."""
        filename = make_path(filename)
        with fits.open(filename, memmap=False) as hdulist:
            try:
                return cls.from_hdulist(hdulist, hdu=hdu)
            except KeyError:
                raise ValueError(
                    f"File {filename} contains no HDU {hdu!r}\n"
                    f"Available: {[_.name for _ in hdulist]}"
                )

    def to_table(self):
        """Convert to `~astropy.table.Table` in ARF format.

        Data format specification: :ref:`gadf:ogip-arf`
        """
        table = Table()
        table.meta = {
            "EXTNAME": "SPECRESP",
            "hduclass": "OGIP",
            "hduclas1": "RESPONSE",
            "hduclas2": "SPECRESP",
        }

        energy = self.energy.edges
        table["ENERG_LO"] = energy[:-1]
        table["ENERG_HI"] = energy[1:]
        table["SPECRESP"] = self.evaluate_fill_nan()
        return table

    def to_hdulist(self, name=None, use_sherpa=False):
        """Convert to `~astropy.io.fits.HDUList`."""
        table = self.to_table()

        if use_sherpa:
            table["ENERG_HI"] = table["ENERG_HI"].quantity.to("keV")
            table["ENERG_LO"] = table["ENERG_LO"].quantity.to("keV")
            table["SPECRESP"] = table["SPECRESP"].quantity.to("cm2")

        return fits.HDUList([fits.PrimaryHDU(), fits.BinTableHDU(table, name=name)])

    def write(self, filename, use_sherpa=False, **kwargs):
        """Write to file."""
        filename = make_path(filename)
        self.to_hdulist(use_sherpa=use_sherpa).writeto(filename, **kwargs)

    def evaluate_fill_nan(self, **kwargs):
        """Modified evaluate function.

        Calls :func:`gammapy.utils.nddata.NDDataArray.evaluate` and replaces
        possible nan values. Below the finite range the effective area is set
        to zero and above to value of the last valid note. This is needed since
        other codes, e.g. sherpa, don't like nan values in FITS files. Make
        sure that the replacement happens outside of the energy range, where
        the `~gammapy.irf.EffectiveAreaTable` is used.
        """
        retval = self.data.evaluate(**kwargs)
        idx = np.where(np.isfinite(retval))[0]
        retval[np.arange(idx[0])] = 0
        retval[np.arange(idx[-1], len(retval))] = retval[idx[-1]]
        return retval

    @property
    def max_area(self):
        """Maximum effective area."""
        cleaned_data = self.data.data[np.where(~np.isnan(self.data.data))]
        return cleaned_data.max()

    def find_energy(self, aeff, emin=None, emax=None):
        """Find energy for a given effective area.

        In case the solution is not unique, provide the `emin` or `emax` arguments
        to limit the solution to the given range. By default the peak energy of the
        effective area is chosen as `emax`.

        Parameters
        ----------
        aeff : `~astropy.units.Quantity`
            Effective area value
        emin : `~astropy.units.Quantity`
            Lower bracket value in case solution is not unique.
        emax : `~astropy.units.Quantity`
            Upper bracket value in case solution is not unique.

        Returns
        -------
        energy : `~astropy.units.Quantity`
            Energy corresponding to the given aeff.
        """
        from gammapy.modeling.models import TemplateSpectralModel

        energy = self.energy.center

        if emin is None:
            emin = energy[0]
        if emax is None:
            # use the peak effective area as a default for the energy maximum
            emax = energy[np.argmax(self.data.data)]

        aeff_spectrum = TemplateSpectralModel(energy, self.data.data)
        return aeff_spectrum.inverse(aeff, emin=emin, emax=emax)

class EffectiveAreaTable2D:
    """2D effective area table.

    Data format specification: :ref:`gadf:aeff_2d`

    Parameters
    ----------
    energy_lo, energy_hi : `~astropy.units.Quantity`
        Energy binning
    offset_lo, offset_hi : `~astropy.units.Quantity`
        Field of view offset angle.
    data : `~astropy.units.Quantity`
        Effective area

    Examples
    --------
    Here's an example you can use to learn about this class:

    >>> from gammapy.irf import EffectiveAreaTable2D
    >>> filename = '$GAMMAPY_DATA/cta-1dc/caldb/data/cta/1dc/bcf/South_z20_50h/irf_file.fits'
    >>> aeff = EffectiveAreaTable2D.read(filename, hdu='EFFECTIVE AREA')
    >>> print(aeff)
    EffectiveAreaTable2D
    NDDataArray summary info
    energy         : size =    42, min =  0.014 TeV, max = 177.828 TeV
    offset         : size =     6, min =  0.500 deg, max =  5.500 deg
    Data           : size =   252, min =  0.000 m2, max = 5371581.000 m2

    Here's another one, created from scratch, without reading a file:

    >>> from gammapy.irf import EffectiveAreaTable2D
    >>> import astropy.units as u
    >>> import numpy as np
    >>> energy = np.logspace(0,1,11) * u.TeV
    >>> offset = np.linspace(0,1,4) * u.deg
    >>> data = np.ones(shape=(10,3)) * u.cm * u.cm
    >>> aeff = EffectiveAreaTable2D(energy_lo=energy[:-1], energy_hi=energy[1:], offset_lo=offset[:-1],
    >>>                             offset_hi=offset[1:], data= data)
    >>> print(aeff)
    Data array summary info
    energy         : size =    11, min =  1.000 TeV, max = 10.000 TeV
    offset         : size =     4, min =  0.000 deg, max =  1.000 deg
    Data           : size =    30, min =  1.000 cm2, max =  1.000 cm2
    """

    default_interp_kwargs = dict(bounds_error=False, fill_value=None)
    """Default Interpolation kwargs for `~NDDataArray`. Extrapolate."""

    def __init__(
        self,
        energy_lo,
        energy_hi,
        offset_lo,
        offset_hi,
        data,
        meta=None,
        interp_kwargs=None,
    ):

        if interp_kwargs is None:
            interp_kwargs = self.default_interp_kwargs

        e_edges = edges_from_lo_hi(energy_lo, energy_hi)
        energy_axis = MapAxis.from_edges(e_edges, interp="log", name="energy")

        # TODO: for some reason the H.E.S.S. DL3 files contain the same values for offset_hi and offset_lo
        if np.allclose(offset_lo.to_value("deg"), offset_hi.to_value("deg")):
            offset_axis = MapAxis.from_nodes(offset_lo, interp="lin", name="offset")
        else:
            offset_edges = edges_from_lo_hi(offset_lo, offset_hi)
            offset_axis = MapAxis.from_edges(offset_edges, interp="lin", name="offset")

        self.data = NDDataArray(
            axes=[energy_axis, offset_axis], data=data, interp_kwargs=interp_kwargs
        )
        self.meta = meta or {}

    def __str__(self):
        ss = self.__class__.__name__
        ss += f"\n{self.data}"
        return ss

    @property
    def low_threshold(self):
        """Low energy threshold"""
        return self.meta["LO_THRES"] * u.TeV

    @property
    def high_threshold(self):
        """High energy threshold"""
        return self.meta["HI_THRES"] * u.TeV

    @classmethod
    def from_table(cls, table):
        """Read from `~astropy.table.Table`."""
        return cls(
            energy_lo=table["ENERG_LO"].quantity[0],
            energy_hi=table["ENERG_HI"].quantity[0],
            offset_lo=table["THETA_LO"].quantity[0],
            offset_hi=table["THETA_HI"].quantity[0],
            data=table["EFFAREA"].quantity[0].transpose(),
            meta=table.meta,
        )

    @classmethod
    def from_hdulist(cls, hdulist, hdu="EFFECTIVE AREA"):
        """Create from `~astropy.io.fits.HDUList`."""
        return cls.from_table(Table.read(hdulist[hdu]))

    @classmethod
    def read(cls, filename, hdu="EFFECTIVE AREA"):
        """Read from file."""
        with fits.open(make_path(filename), memmap=False) as hdulist:
            return cls.from_hdulist(hdulist, hdu=hdu)

    def to_effective_area_table(self, offset, energy=None):
        """Evaluate at a given offset and return `~gammapy.irf.EffectiveAreaTable`.

        Parameters
        ----------
        offset : `~astropy.coordinates.Angle`
            Offset
        energy : `~astropy.units.Quantity`
            Energy axis bin edges
        """
        if energy is None:
            energy = self.data.axis("energy").edges

        area = self.data.evaluate(
            offset=offset, energy=MapAxis.from_edges(energy, interp="log").center
        )

        return EffectiveAreaTable(
            energy_lo=energy[:-1], energy_hi=energy[1:], data=area
        )

    def plot_energy_dependence(self, ax=None, offset=None, energy=None, **kwargs):
        """Plot effective area versus energy for a given offset.

        Parameters
        ----------
        ax : `~matplotlib.axes.Axes`, optional
            Axis
        offset : `~astropy.coordinates.Angle`
            Offset
        energy : `~astropy.units.Quantity`
            Energy axis
        kwargs : dict
            Forwarded tp plt.plot()

        Returns
        -------
        ax : `~matplotlib.axes.Axes`
            Axis
        """
        import matplotlib.pyplot as plt

        ax = plt.gca() if ax is None else ax

        if offset is None:
            off_min, off_max = self.data.axis("offset").center[[0, -1]]
            offset = np.linspace(off_min.value, off_max.value, 4) * off_min.unit

        if energy is None:
            energy = self.data.axis("energy").center

        for off in offset:
            area = self.data.evaluate(offset=off, energy=energy)
            label = f"offset = {off:.1f}"
            ax.plot(energy, area.value, label=label, **kwargs)

        ax.set_xscale("log")
        ax.set_xlabel(f"Energy [{self.data.axis('energy').unit}]")
        ax.set_ylabel(f"Effective Area [{self.data.data.unit}]")
        ax.set_xlim(min(energy.value), max(energy.value))
        ax.legend(loc="upper left")

        return ax

    def plot_offset_dependence(self, ax=None, offset=None, energy=None, **kwargs):
        """Plot effective area versus offset for a given energy.

        Parameters
        ----------
        ax : `~matplotlib.axes.Axes`, optional
            Axis
        offset : `~astropy.coordinates.Angle`
            Offset axis
        energy : `~astropy.units.Quantity`
            Energy

        Returns
        -------
        ax : `~matplotlib.axes.Axes`
            Axis
        """
        import matplotlib.pyplot as plt

        ax = plt.gca() if ax is None else ax

        if energy is None:
            e_min, e_max = np.log10(self.data.axis("energy").center.value[[0, -1]])
            energy = np.logspace(e_min, e_max, 4) * self.data.axis("energy").unit

        if offset is None:
            offset = self.data.axis("offset").center

        for ee in energy:
            area = self.data.evaluate(offset=offset, energy=ee)
            area /= np.nanmax(area)
            if np.isnan(area).all():
                continue
            label = f"energy = {ee:.1f}"
            ax.plot(offset, area, label=label, **kwargs)

        ax.set_ylim(0, 1.1)
        ax.set_xlabel(f"Offset ({self.data.axis('offset').unit})")
        ax.set_ylabel("Relative Effective Area")
        ax.legend(loc="best")

        return ax

    def plot(self, ax=None, add_cbar=True, **kwargs):
        """Plot effective area image."""
        import matplotlib.pyplot as plt

        ax = plt.gca() if ax is None else ax

        energy = self.data.axis("energy").edges
        offset = self.data.axis("offset").edges
        aeff = self.data.evaluate(offset=offset, energy=energy[:, np.newaxis])

        vmin, vmax = np.nanmin(aeff.value), np.nanmax(aeff.value)

        kwargs.setdefault("cmap", "GnBu")
        kwargs.setdefault("edgecolors", "face")
        kwargs.setdefault("vmin", vmin)
        kwargs.setdefault("vmax", vmax)

        caxes = ax.pcolormesh(energy.value, offset.value, aeff.value.T, **kwargs)

        ax.set_xscale("log")
        ax.set_ylabel(f"Offset ({offset.unit})")
        ax.set_xlabel(f"Energy ({energy.unit})")

        xmin, xmax = energy.value.min(), energy.value.max()
        ax.set_xlim(xmin, xmax)

        if add_cbar:
            label = f"Effective Area ({aeff.unit})"
            ax.figure.colorbar(caxes, ax=ax, label=label)

        return ax

    def peek(self, figsize=(15, 5)):
        """Quick-look summary plots."""
        import matplotlib.pyplot as plt

        fig, axes = plt.subplots(nrows=1, ncols=3, figsize=figsize)
        self.plot(ax=axes[2])
        self.plot_energy_dependence(ax=axes[0])
        self.plot_offset_dependence(ax=axes[1])
        plt.tight_layout()

    def to_table(self):
        """Convert to `~astropy.table.Table`."""
        meta = self.meta.copy()

        energy = self.data.axis("energy").edges
        theta = self.data.axis("offset").edges

        table = Table(meta=meta)
        table["ENERG_LO"] = energy[:-1][np.newaxis]
        table["ENERG_HI"] = energy[1:][np.newaxis]
        table["THETA_LO"] = theta[:-1][np.newaxis]
        table["THETA_HI"] = theta[1:][np.newaxis]
        table["EFFAREA"] = self.data.data.T[np.newaxis]
        return table

    def to_fits(self, name="EFFECTIVE AREA"):
        """Convert to `~astropy.io.fits.BinTableHDU`."""
        return fits.BinTableHDU(self.to_table(), name=name)

import numpy,import astropy,from astropy


# Copyright 2021 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================

"""
Implement octree.
"""

import numpy as np

def int2node_data(value):
    data = np.zeros(8, dtype=np.int)
    value_vector = np.repeat(value, 8)
    tmp = np.floor(value_vector / [1, 2, 4, 8, 16, 32, 64, 128])
    data = np.mod(tmp, 2)
    return data.astype(int)

class OctreeBranchNode:
    def __init__(self, depth, parent, position, size, octant):
        """
        Here we use number to simply eight branches
        + here means greater then the center, vice versa
        branch: 0 1 2 3 4 5 6 7
        x:      - - - - + + + +
        y:      - - + + - - + +
        z:      - + - + - + - +

        Parameters
        ----------
        position: ndarray
            center of voxel/node
        """
        self.attribute = "branch"
        self.depth = depth
        self.size = size
        self.position = position
        self.parent = parent
        self.data = np.zeros(8, dtype=np.int)
        self.branches = [None, None, None, None, None, None, None, None]
        self.octant = octant

    def __str__(self):
        return u"position: {0}, size: {1}, depth: {2} parent: {3}, data: {4}".format(
            self.position, self.size, self.depth, self.parent, self.data)

class OctreeLeafNode:

    def __init__(self, depth, parent, position_gt, size):
        """
        leafnode do not need branch and data
        Different with position in the branch, position_gt here is not the center of voxel bur ground truth xyz
        """
        self.attribute = "leaf"
        self.depth = depth
        self.size = size
        self.position_center = None
        self.position_gt = position_gt
        self.parent = parent

class Octree:

    def __init__(self, max_range, precision, origin=(0, 0, 0)):
        """
        Parameters
        ----------
        max_range:
            the size of whole octree
        origin:
            the position of octree root node
        leaf_num:
            number of lead nodes, aka the number of points after converting into octree
        """
        self.attribute = "root"
        self.size = max_range
        self.precision = precision
        self.depth = 0
        self.max_depth = int(np.log2(max_range / precision)) + 1
        self.position = origin
        self.data = np.zeros(8, dtype=np.int)
        self.branches = [None, None, None, None, None, None, None, None]
        self.leaf_num = 0

    @staticmethod
    def find_branch(root, position):
        """
        helper function
        returns an index corresponding to a branch
        pointing in the direction we want to go
        """
        index = 0
        if position[0] >= root.position[0]:
            index |= 4
        if position[1] >= root.position[1]:
            index |= 2
        if position[2] >= root.position[2]:
            index |= 1
        return index

    def insert_node(self, root, size, parent, position, depth):
        if depth == self.max_depth:
            if root is None:
                self.leaf_num += 1
            return OctreeLeafNode(depth, parent, position, size)

        if depth < self.max_depth:
            branch = self.find_branch(root, position)
            branch_size = root.size / 2
            root.data[branch] = 1

            if (root.branches[branch] is None) and (depth != self.max_depth - 1):
                pos = root.position
                offset = size / 2
                new_center = (0, 0, 0)
                if branch == 0:
                    new_center = (pos[0] - offset, pos[1] - offset, pos[2] - offset)
                elif branch == 1:
                    new_center = (pos[0] - offset, pos[1] - offset, pos[2] + offset)
                elif branch == 2:
                    new_center = (pos[0] - offset, pos[1] + offset, pos[2] - offset)
                elif branch == 3:
                    new_center = (pos[0] - offset, pos[1] + offset, pos[2] + offset)
                elif branch == 4:
                    new_center = (pos[0] + offset, pos[1] - offset, pos[2] - offset)
                elif branch == 5:
                    new_center = (pos[0] + offset, pos[1] - offset, pos[2] + offset)
                elif branch == 6:
                    new_center = (pos[0] + offset, pos[1] + offset, pos[2] - offset)
                elif branch == 7:
                    new_center = (pos[0] + offset, pos[1] + offset, pos[2] + offset)
                root.branches[branch] = OctreeBranchNode(depth + 1, root, new_center, branch_size, branch)

            root.branches[branch] = self.insert_node(root.branches[branch], branch_size, root, position, depth + 1)
        return root

    def serialize_depth_first(self):
        def extract_data(node):
            if node and node.attribute != "leaf":
                values.append(node_data2int(node.data))
                # vals.append(node.data)
                for branch_id in range(8):
                    extract_data(node.branches[branch_id])

        values = []
        extract_data(self)
        return values

def deserialize_depth_first(values, max_depth, octree):
    def reconstruct_tree(root):
        value = next(values)
        node_data = int2node_data(value)
        root.data = node_data
        for branch in range(8):
            if node_data[branch] == 0:
                root.branches[branch] = None
            else:
                depth = root.depth
                pos = root.position
                offset = root.size / 2
                new_center = (0, 0, 0)

                if branch == 0:
                    new_center = (pos[0] - offset, pos[1] - offset, pos[2] - offset)
                elif branch == 1:
                    new_center = (pos[0] - offset, pos[1] - offset, pos[2] + offset)
                elif branch == 2:
                    new_center = (pos[0] - offset, pos[1] + offset, pos[2] - offset)
                elif branch == 3:
                    new_center = (pos[0] - offset, pos[1] + offset, pos[2] + offset)
                elif branch == 4:
                    new_center = (pos[0] + offset, pos[1] - offset, pos[2] - offset)
                elif branch == 5:
                    new_center = (pos[0] + offset, pos[1] - offset, pos[2] + offset)
                elif branch == 6:
                    new_center = (pos[0] + offset, pos[1] + offset, pos[2] - offset)
                elif branch == 7:
                    new_center = (pos[0] + offset, pos[1] + offset, pos[2] + offset)

                if depth < max_depth - 1:
                    node = OctreeBranchNode(depth + 1, root, new_center, offset, branch)
                    root.branches[branch] = node
                    reconstruct_tree(root.branches[branch])
                else:
                    node = OctreeLeafNode(depth + 1, root, 0, offset)
                    node.position_center = new_center
                    recon_points.append(new_center)
                    root.branches[branch] = node
        return root
    recon_points = []
    return reconstruct_tree(octree), recon_points


# Copyright 2020 Makani Technologies LLC
#
# 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.

"""Scoring functions relating to aerodynamics."""

from makani.analysis.aero import aero_ssam
from makani.analysis.aero import apparent_wind_util
from makani.analysis.log_analysis import loop_averager
from makani.control import system_types
from makani.lib.python import c_helpers
from makani.lib.python.batch_sim import scoring_functions
from makani.lib.python.h5_utils import numpy_utils
from makani.system import labels as system_labels

import numpy as np
from scipy import interpolate
import scoring_functions_util as scoring_util

_FLAP_LABEL_HELPER = c_helpers.EnumHelper('FlapLabel', system_labels,
                                          prefix='kFlap')
_WING_MODEL_HELPER = c_helpers.EnumHelper('WingModel', system_types)
_WING_SERIAL_HELPER = c_helpers.EnumHelper('WingSerial', system_types)

class AirspeedMaxScoringFunction(
    scoring_functions.SingleSidedLimitScoringFunction):
  """Tests if the airspeed falls outside of acceptable limits."""

  def __init__(self, good_limit, bad_limit, severity):
    super(AirspeedMaxScoringFunction, self).__init__(
        'Max Airspeed', 'm/s', good_limit, bad_limit, severity)

  def GetSystemLabels(self):
    return ['controls']

  def GetValue(self, output):
    return output['airspeed_max']

  def GetOutput(self, timeseries):
    return {'airspeed_max': np.max(timeseries['airspeed'])}

  def GetTimeSeries(self, params, sim, control):
    airspeed = self._SelectTelemetry(sim, control, 'airspeed')
    return {'airspeed': airspeed}

class AirspeedMinScoringFunction(
    scoring_functions.SingleSidedLimitScoringFunction):
  """Tests if the airspeed falls outside of acceptable limits."""

  def __init__(self, good_limit, bad_limit, severity):
    super(AirspeedMinScoringFunction, self).__init__(
        'Min Airspeed', 'm/s', good_limit, bad_limit, severity)

  def GetSystemLabels(self):
    return ['controls']

  def GetValue(self, output):
    return output['airspeed_min']

  def GetOutput(self, timeseries):
    return {'airspeed_min': np.min(timeseries['airspeed'])}

  def GetTimeSeries(self, params, sim, control):
    airspeed = self._SelectTelemetry(sim, control, 'airspeed')
    return {'airspeed': airspeed}

class MainWingAlphaScoringFunction(
    scoring_functions.DoubleSidedLimitScoringFunction):
  """Tests if angle of attack on the main wing exceeds a limit."""

  def __init__(self, bad_lower_limit, good_lower_limit, good_upper_limit,
               bad_upper_limit, severity, airspeed_threshold=1.0,
               steady_flight=False):
    super(MainWingAlphaScoringFunction, self).__init__(
        'Main Wing SSAM AoA%s' % (' (w/o initial transients)' if steady_flight
                                  else ''),
        'deg', bad_lower_limit, good_lower_limit,
        good_upper_limit, bad_upper_limit, severity)
    assert airspeed_threshold > 0.0
    self._airspeed_threshold = airspeed_threshold

  def GetSystemLabels(self):
    return ['aero']

  def GetValue(self, output):
    return np.array([output['wing_alpha_min'], output['wing_alpha_max']])

  def GetFailureCount(self, timeseries):
    if timeseries is None:
      return None
    # Confirm that the arrays have the same size.
    assert len(timeseries['alphas_min']) == len(timeseries['alphas_max'])
    # Confirm that at any index, alphas_max is going to be greater than
    # alphas_min.
    assert np.all(timeseries['alphas_max'] >= timeseries['alphas_min'])
    return np.logical_or(
        timeseries['alphas_min'] < self._bad_lower_limit,
        timeseries['alphas_max'] > self._bad_upper_limit).sum()

  def GetIndexOfFirstFailure(self, timeseries):
    if timeseries is None:
      return None
    if self.GetFailureCount(timeseries) == 0:
      return -1
    else:
      return np.argmax(
          np.logical_or(timeseries['alphas_min'] < self._bad_lower_limit,
                        timeseries['alphas_max'] > self._bad_upper_limit))

  def GetOutput(self, timeseries):
    return {
        'wing_alpha_min': np.min(timeseries['alphas_min']),
        'wing_alpha_max': np.max(timeseries['alphas_max'])
    }

  def GetTimeSeries(self, params, sim, control):
    airspeeds, angular_rates, app_wind_b = self._SelectTelemetry(
        sim, control, ['airspeed', 'body_rates', 'apparent_wind_vector'])

    # Converts the indices into an (n,) sized array for 2D array masking.
    data_indices = np.reshape(
        np.argwhere(airspeeds > self._airspeed_threshold), -1)

    if data_indices.size == 0:
      wing_alphas_max = np.array([float('nan')])
      wing_alphas_min = np.array([float('nan')])

    else:
      # Mask the body rates for telemetry that crosses the threshold.
      omega_b = numpy_utils.Vec3ToArray(angular_rates)[data_indices]

      # Mask the cartesian wind for telemetry that crosses the threshold.
      wind_b = numpy_utils.Vec3ToArray(app_wind_b)[data_indices]

      # Compute the kinematic-based local values of alpha.
      # TODO: Wing_model should be mapped to enumerate wing models.
      wing_model = _WING_MODEL_HELPER.ShortName(
          int(params['system_params']['wing_model'][0]))
      wing_serial = _WING_SERIAL_HELPER.Name(int(
          params['system_params']['wing_serial'][0]))
      ssam = aero_ssam.SSAMModel(wing_model, wing_serial)
      wing_alphas_deg = ssam.GetMainWingAlphas(omega_b, wind_b)

      # Provide telemetry of maximum alphas anywhere along the main wing.
      # As per GetMainWingAlphas the expected size of wing_alphas_deg is (n, m)
      # where m is the number of wing panels and n is the number of elements in
      # the time series.
      wing_alphas_max = np.amax(wing_alphas_deg, axis=1)
      wing_alphas_min = np.amin(wing_alphas_deg, axis=1)

    return {'alphas_max': wing_alphas_max, 'alphas_min': wing_alphas_min}

class AlphaDegScoringFunction(
    scoring_functions.DoubleSidedLimitScoringFunction):
  """Tests if the angle-of-attack exceeds a limit when airspeed is high."""

  def __init__(self, bad_lower_limit, good_lower_limit, good_upper_limit,
               bad_upper_limit, severity, airspeed_threshold=1.0,
               steady_flight=False):
    super(AlphaDegScoringFunction, self).__init__(
        'Angle-of-attack%s' % (' (w/o initial transients)' if steady_flight
                               else ''),
        'deg', bad_lower_limit, good_lower_limit,
        good_upper_limit, bad_upper_limit, severity)
    assert airspeed_threshold > 0.0
    self._airspeed_threshold = airspeed_threshold

  def GetSystemLabels(self):
    return ['aero']

  def GetValue(self, output):
    return np.array([output['alpha_min'], output['alpha_max']])

  def GetOutput(self, timeseries):
    alpha = timeseries['alpha']
    return {
        'alpha_max': np.max(alpha),
        'alpha_min': np.min(alpha)
    }

  def GetTimeSeries(self, params, sim, control):
    airspeeds, alpha, time = self._SelectTelemetry(
        sim, control, ['airspeed', 'alpha', 'time'])
    data_indices = np.argwhere(airspeeds > self._airspeed_threshold)
    if data_indices.size == 0:
      alpha = np.array([float('nan')])
    else:
      chord = params['system_params']['wing']['c']
      # Choose 5 main wing flow-over time constants, 2 sigma for peak rejection
      alpha = scoring_util.FilterByWindowAveraging(
          time[data_indices], alpha[data_indices], airspeeds[data_indices],
          chord, num_tau=5, num_sigma=2)

    return {'alpha': np.rad2deg(alpha)}

class MinAlphaDegScoringFunction(
    scoring_functions.DoubleSidedLimitScoringFunction):
  """Tests if the angle-of-attack exceeds a limit when airspeed is high."""

  def __init__(self, bad_lower_limit, good_lower_limit, good_upper_limit,
               bad_upper_limit, severity, airspeed_threshold=1.0):
    super(MinAlphaDegScoringFunction, self).__init__(
        'Min. Angle-of-attack', 'deg', bad_lower_limit, good_lower_limit,
        good_upper_limit, bad_upper_limit, severity)
    assert airspeed_threshold > 0.0
    self._airspeed_threshold = airspeed_threshold

  def GetSystemLabels(self):
    return ['controls']

  def GetValue(self, output):
    return output['min_alpha']

  def GetOutput(self, timeseries):
    return {'min_alpha': np.min(timeseries['alpha'])}

  def GetTimeSeries(self, params, sim, control):
    # TODO: Review if this scoring function is working as
    # intended. Update/remove as needed.
    airspeeds, alpha, time = self._SelectTelemetry(
        sim, control, ['airspeed', 'alpha', 'time'])
    data_indices = np.argwhere(airspeeds > self._airspeed_threshold)
    if data_indices.size == 0:
      alpha = np.array([float('nan')])
    else:
      chord = params['system_params']['wing']['c']
      # Choose 5 main wing flow-over time constants, 2 sigma for peak rejection
      alpha = scoring_util.FilterByWindowAveraging(
          time[data_indices], alpha[data_indices], airspeeds[data_indices],
          chord, num_tau=5, num_sigma=2)

    return {'alpha': np.rad2deg(alpha)}

class BetaDegScoringFunction(scoring_functions.DoubleSidedLimitScoringFunction):
  """Tests if the side-slip exceeds a limit when airspeed is high."""

  def __init__(self, bad_lower_limit, good_lower_limit,
               good_upper_limit, bad_upper_limit, severity,
               airspeed_threshold=1.0, steady_flight=False):
    super(BetaDegScoringFunction, self).__init__(
        'Side-slip%s' % (' (w/o initial transients)' if steady_flight else ''),
        'deg', bad_lower_limit, good_lower_limit,
        good_upper_limit, bad_upper_limit, severity)
    assert airspeed_threshold > 0.0
    self._airspeed_threshold = airspeed_threshold

  def GetSystemLabels(self):
    return ['aero']

  def GetValue(self, output):
    return np.array([output['beta_min'], output['beta_max']])

  def GetOutput(self, timeseries):
    beta = timeseries['beta']
    return {
        'beta_max': np.max(beta),
        'beta_min': np.min(beta)
    }

  def GetTimeSeries(self, params, sim, control):
    airspeeds, beta, time = self._SelectTelemetry(
        sim, control, ['airspeed', 'beta', 'time'])
    data_indices = np.argwhere(airspeeds > self._airspeed_threshold)
    if data_indices.size == 0:
      beta = np.array([float('nan')])
    else:
      chord = params['system_params']['wing']['c']
      # Choose 5 main wing flow-over time constants, 2 sigma for peak rejection
      beta = scoring_util.FilterByWindowAveraging(
          time[data_indices], beta[data_indices], airspeeds[data_indices],
          chord, num_tau=5, num_sigma=2)

    return {'beta': np.rad2deg(beta)}

class AlphaDegErrorScoringFunction(
    scoring_functions.SingleSidedLimitScoringFunction):
  """Tests if the angle-of-attack error in crosswind flight is small enough."""

  def __init__(self, good_limit, bad_limit, severity,
               airspeed_threshold=1.0, steady_flight=False):
    super(AlphaDegErrorScoringFunction, self).__init__(
        'Angle-of-attack Error%s' % (' (w/o initial transients)'
                                     if steady_flight else ''),
        'deg', good_limit, bad_limit, severity)
    assert airspeed_threshold > 0.0
    self._airspeed_threshold = airspeed_threshold
    self.SetSourcePriority(['control'])

  def GetSystemLabels(self):
    return ['controls']

  def GetValue(self, output):
    return output['alpha_error_max']

  def GetOutput(self, timeseries):
    return {'alpha_error_max': np.max(timeseries['alpha_error'])}

  def GetTimeSeries(self, params, sim, control):
    airspeeds, alphas, alpha_cmds, time = self._SelectTelemetry(
        sim, control, ['airspeed', 'alpha', 'alpha_cmd', 'time'])

    data_indices = np.argwhere(airspeeds > self._airspeed_threshold)
    if data_indices.size == 0:
      alpha_error = np.array([float('nan')])
    else:
      chord = params['system_params']['wing']['c']
      # Choose 5 main wing flow-over time constants, 2 sigma for peak rejection
      alphas = scoring_util.FilterByWindowAveraging(
          time[data_indices], alphas[data_indices], airspeeds[data_indices],
          chord, num_tau=5, num_sigma=2)
      alpha_error = np.fabs(alphas - alpha_cmds[data_indices])

    return {'alpha_error': np.rad2deg(alpha_error)}

class BetaDegErrorScoringFunction(
    scoring_functions.SingleSidedLimitScoringFunction):
  """Tests if the sideslip angle error in crosswind flight is small enough."""

  def __init__(self, good_limit, bad_limit, severity,
               airspeed_threshold=1.0, steady_flight=False):
    super(BetaDegErrorScoringFunction, self).__init__(
        'Sideslip Error%s' % (' (w/o initial transients)'
                              if steady_flight else ''),
        'deg', good_limit, bad_limit, severity)
    assert airspeed_threshold > 0.0
    self._airspeed_threshold = airspeed_threshold
    self.SetSourcePriority(['control'])

  def GetSystemLabels(self):
    return ['controls']

  def GetValue(self, output):
    return output['beta_error_max']

  def GetOutput(self, timeseries):
    return {'beta_error_max': np.max(timeseries['beta_error'])}

  def GetTimeSeries(self, params, sim, control):
    airspeeds, betas, beta_cmds, time = self._SelectTelemetry(
        sim, control, ['airspeed', 'beta', 'beta_cmd', 'time'])

    data_indices = np.argwhere(airspeeds > self._airspeed_threshold)
    if data_indices.size == 0:
      beta_error = np.array([float('nan')])
    else:
      chord = params['system_params']['wing']['c']
      # Choose 5 main wing flow-over time constants, 2 sigma for peak rejection
      betas = scoring_util.FilterByWindowAveraging(
          time[data_indices], betas[data_indices], airspeeds[data_indices],
          chord, num_tau=5, num_sigma=2)
      beta_error = np.fabs(betas - beta_cmds[data_indices])

    return {'beta_error': np.rad2deg(beta_error)}

# TODO: Alpha, Beta and Airspeed RMS error scoring functions
# should be evaluated on per loop basis. The GetTimeseries method should be like
# the one in SurfaceSaturationScoringFunction.
class AlphaRmsScoringFunction(
    scoring_functions.SingleSidedLimitScoringFunction):
  """RMS angle-of-attack error."""

  def __init__(self, good_limit, bad_limit, severity):
    super(AlphaRmsScoringFunction, self).__init__(
        'AoA RMS', 'deg', good_limit, bad_limit, severity)
    self.SetSourcePriority(['control'])

  def GetSystemLabels(self):
    return ['experimental', 'controls']

  def GetValue(self, output):
    return output['alpha_rms']

  def GetOutput(self, timeseries):
    return {'alpha_rms': np.rad2deg(timeseries['alpha_rms'])}

  def GetTimeSeries(self, params, sim, control):
    alphas, alpha_cmds = self._SelectTelemetry(
        sim, control, ['alpha', 'alpha_cmd'])
    return {
        'alpha_rms': np.mean((alphas - alpha_cmds)**2.0)**0.5
        }

class BetaRmsScoringFunction(
    scoring_functions.SingleSidedLimitScoringFunction):
  """RMS side-slip error."""

  def __init__(self, good_limit, bad_limit, severity):
    super(BetaRmsScoringFunction, self).__init__(
        'Side-slip RMS', 'deg', good_limit, bad_limit, severity)
    self.SetSourcePriority(['control'])

  def GetSystemLabels(self):
    return ['experimental', 'controls']

  def GetValue(self, output):
    return output['beta_rms']

  def GetOutput(self, timeseries):
    return {'beta_rms': np.rad2deg(timeseries['beta_rms'])}

  def GetTimeSeries(self, params, sim, control):
    betas, beta_cmds = self._SelectTelemetry(
        sim, control, ['beta', 'beta_cmd'])
    return {
        'beta_rms': np.mean((betas - beta_cmds)**2.0)**0.5
        }

class AirspeedRmsScoringFunction(
    scoring_functions.SingleSidedLimitScoringFunction):
  """RMS airspeed error."""

  def __init__(self, good_limit, bad_limit, severity):
    super(AirspeedRmsScoringFunction, self).__init__(
        'Airspeed RMS', 'm/s', good_limit, bad_limit, severity)
    self.SetSourcePriority(['control'])

  def GetSystemLabels(self):
    return ['experimental', 'controls']

  def GetValue(self, output):
    return output['airspeed_rms']

  def GetOutput(self, timeseries):
    return {'airspeed_rms': timeseries['airspeed_rms']}

  def GetTimeSeries(self, params, sim, control):
    airspeeds, airspeed_cmds = self._SelectTelemetry(
        sim, control, ['airspeed', 'airspeed_cmd'])
    return {
        'airspeed_rms': np.mean((airspeeds - airspeed_cmds)**2.0)**0.5
        }

class SurfaceSaturationScoringFunction(
    scoring_functions.SingleSidedLimitScoringFunction):
  """Tests for the percent time the control surface is saturated.

     If the saturation scoring function is activated during a crosswind flight
     mode, then the characteristic percentage time corresponds to the maximum
     percentage time a surface was saturated during any particular single loop.

     If the saturation scoring function is activated during a non-crosswind
     flight mode, then the characteristic percentage time corresponds to the
     maximum saturation percentage achieved on all the telemetry data for
     that flight mode.
  """

  def __init__(self, flap_labels, good_limit, bad_limit, severity):
    self._flap_indices = [_FLAP_LABEL_HELPER.Value(l) for l in flap_labels]
    super(SurfaceSaturationScoringFunction, self).__init__(
        '%s %% Saturated' % '/'.join(flap_labels), '% time',
        good_limit, bad_limit, severity)

  def GetSystemLabels(self):
    return ['aero', 'controls']

  def GetValue(self, output):
    return output['percent_time_sat']

  def GetOutput(self, timeseries):
    largest_saturation_loop_mask = timeseries['largest_saturation_loop_mask']
    if largest_saturation_loop_mask is None:
      return {'percent_time_sat': float('nan')}
    saturation_counter = np.where(largest_saturation_loop_mask)[0]
    return {
        'percent_time_sat': float(np.size(saturation_counter)) /
                            np.size(largest_saturation_loop_mask) * 100.0
    }

  def GetTimeSeries(self, params, sim, control):
    # "_i" is used to denote the as-imported variable.
    (omega_i, airspeed_i, alpha_i, beta_i, flaps, loop_angles) = (
        self._SelectTelemetry(sim, control, [
            'body_rates', 'airspeed', 'alpha', 'beta', 'flaps', 'loop_angle']))
    if not scoring_util.IsSelectionValid(flaps):
      return {'saturation_mask': None, 'largest_saturation_loop_mask': None}
    deflections = np.rad2deg(flaps[:, self._flap_indices])

    # Observations of actual control surface deflections show that when
    # saturated, they are not exactly at the control limit. Adding or
    # subtracting 0.25 degrees from the limit allows for most saturations
    # to be flagged by the criteria. A4 and A5 are often very close to the
    # upper limit of zero with very small magnitudes of deflection. An offset
    # of only 0.005 degrees is applied to the upper flap limit for flaps A4
    # and A5 to prevent normal operation from being flagged erroneously.
    # The deflection limits of the rudder are altered to account for the
    # wing-fuselage junction loads limit, per b/112267831.
    lower_deflection_limit = (
        np.rad2deg(np.array(params['control_params']['crosswind']['output']
                            ['lower_flap_limits'][0, self._flap_indices]))
        + 0.25)
    if (self._flap_indices == [system_labels.kFlapA4] or
        self._flap_indices == [system_labels.kFlapA5]):
      upper_deflection_limit = (np.rad2deg(np.array(
          params['control_params']['crosswind']['output']
          ['upper_flap_limits'][0, self._flap_indices])) - 0.005)

    elif self._flap_indices == [system_labels.kFlapRud]:
      # Check if body rates data exist. These may not exist if the relevant
      # flight modes do not exist.
      if not scoring_util.IsSelectionValid(omega_i):
        return {'saturation_mask': float('nan')}
      # Position [m] of the vtail aerodynamic center. Only the x-coordinate is
      # relevant here.
      # TODO: Add position of aerodynamic center to params in the h5 log.
      r_vtail = np.array([-7.0, 0.0, 0.0])

      vapp = np.array(airspeed_i)
      alpha = np.array(alpha_i)
      beta = np.array(beta_i)

      omega = np.array([omega_i['x'],
                        omega_i['y'],
                        omega_i['z']])
      # Account for motion of empennage relative to kite origin as it affects
      # apparent wind, alpha, beta.
      v_kite = apparent_wind_util.ApparentWindSphToCart(vapp, alpha, beta).T
      v_omega = np.cross(omega, r_vtail, axis=0)
      vapp, alpha, beta = (
          apparent_wind_util.ApparentWindCartToSph((v_kite + v_omega).T))
      rudder_limits = self._GetRudderLimits(beta, vapp)
      lower_deflection_limit = rudder_limits['lower_limit']
      upper_deflection_limit = rudder_limits['upper_limit']

      # Check that there is not a dimension problem coming out of the rudder
      # deflection limit table. Deflections is expected shape (N,1). The rudder
      # table uses the apparent_wind_util that may create arrays of shape (N,).
      if np.shape(lower_deflection_limit) != np.shape(deflections):
        lower_deflection_limit = np.reshape(lower_deflection_limit,
                                            np.shape(deflections))
      if np.shape(upper_deflection_limit) != np.shape(deflections):
        upper_deflection_limit = np.reshape(upper_deflection_limit,
                                            np.shape(deflections))

    else:
      upper_deflection_limit = (np.rad2deg(np.array(
          params['control_params']['crosswind']['output']
          ['upper_flap_limits'][0, self._flap_indices])) - 0.25)

    # Fraction of a limit at which a surface is considered nearly saturated.
    saturation_fraction = 0.9
    new_upper_limit = saturation_fraction * upper_deflection_limit
    new_lower_limit = saturation_fraction * lower_deflection_limit

    is_saturated_upper = np.greater_equal(deflections, new_upper_limit)
    is_saturated_lower = np.less_equal(deflections, new_lower_limit)
    saturation_mask = np.logical_or(is_saturated_upper, is_saturated_lower)
    largest_saturation_loop_mask = saturation_mask

    # Identify the number of loops in the data and cycle over them to find which
    # loop has the highest percentage of saturation. If there are no loops then
    # simply take all the data as the scoring function is likely operating in a
    # non-crosswind mode, e.g. hover.
    endloop_indices = loop_averager.GetEndLoopIndices(loop_angles)
    if np.size(endloop_indices):
      # Go through each independent loop to find the mask array that contains
      # the largest amount of masked values and return that array for scoring.
      for i, endloop_indx in enumerate(endloop_indices):
        if i == 0:
          start_loop_indx = 0
          largest_pct_saturated = 0.0
          largest_saturation_loop_mask = saturation_mask[start_loop_indx:
                                                         endloop_indx+1]
        else:
          start_loop_indx = endloop_indices[i-1]

        # Figure out how long the surface has been saturated during this loop.
        this_loop_mask = saturation_mask[start_loop_indx:endloop_indx + 1]

        saturation_counter = np.where(this_loop_mask)[0]
        percent_saturated = (float(np.size(saturation_counter)) /
                             np.size(this_loop_mask) * 100.0)

        if percent_saturated > largest_pct_saturated:
          largest_saturation_loop_mask = this_loop_mask
          largest_pct_saturated = percent_saturated

    return {'saturation_mask': saturation_mask,
            'largest_saturation_loop_mask': largest_saturation_loop_mask}

  def GetFailureCount(self, timeseries):
    """Returns the counts of True in the `saturation_mask` array."""
    if timeseries is None:
      return None
    return np.sum(timeseries['saturation_mask'])

  def GetIndexOfFirstFailure(self, timeseries):
    """Returns the index of the first True in the `saturation_mask` array."""
    if timeseries is None:
      return None
    if self.GetFailureCount(timeseries) == 0:
      return -1
    else:
      return np.where(timeseries['saturation_mask'])[0][0]

  def _GetRudderLimits(self, beta, airspeed):
    # Tables below are copied from https://goo.gl/BmyhAe. See spreadsheet for
    # description of how the tables are derived.
    betas = np.array([-20.0, -10.0, 0.0, 10.0, 20.0])
    airspeeds = np.array([50.0, 55.0, 60.0, 65.0, 70.0, 75.0,
                          80.0, 85.0, 90.0, 95.0, 100.0])
    rudder_limit_lower = np.array([[-22.00, -22.00, -22.00, -22.00, -22.00],
                                   [-22.00, -22.00, -22.00, -22.00, -22.00],
                                   [-22.00, -22.00, -22.00, -22.00, -16.42],
                                   [-22.00, -22.00, -22.00, -22.00, -10.57],
                                   [-22.00, -22.00, -22.00, -19.54, -5.92],
                                   [-22.00, -22.00, -22.00, -15.72, -2.16],
                                   [-22.00, -22.00, -22.00, -12.59, 0.91],
                                   [-22.00, -22.00, -22.00, -10.00, 3.45],
                                   [-22.00, -22.00, -21.31, -7.83, 5.58],
                                   [-22.00, -22.00, -19.49, -5.99, 7.39],
                                   [-22.00, -22.00, -17.94, -4.43, 8.93]])
    rudder_limit_upper = np.array([[22.00, 22.00, 22.00, 22.00, 22.00],
                                   [17.67, 22.00, 22.00, 22.00, 22.00],
                                   [9.87, 22.00, 22.00, 22.00, 22.00],
                                   [3.80, 17.64, 22.00, 22.00, 22.00],
                                   [-1.02, 12.83, 22.00, 22.00, 22.00],
                                   [-4.91, 8.94, 21.97, 22.00, 22.00],
                                   [-8.09, 5.76, 18.88, 22.00, 22.00],
                                   [-10.73, 3.12, 16.32, 22.00, 22.00],
                                   [-12.94, 0.91, 14.17, 22.00, 22.00],
                                   [-14.81, -0.96, 12.35, 22.00, 22.00],
                                   [-16.40, -2.55, 10.80, 22.00, 22.00]])

    lower_limit_lookup = interpolate.interp2d(
        betas, airspeeds, rudder_limit_lower, kind='linear')
    upper_limit_lookup = interpolate.interp2d(
        betas, airspeeds, rudder_limit_upper, kind='linear')

    lower_limit = np.zeros(len(beta))
    upper_limit = np.zeros(len(beta))
    for ii in range(len(beta)):
      lower_limit[ii] = lower_limit_lookup(
          np.rad2deg(beta[ii]), airspeed[ii]) + 0.25
      upper_limit[ii] = upper_limit_lookup(
          np.rad2deg(beta[ii]), airspeed[ii]) - 0.25

    return {
        'lower_limit': lower_limit,
        'upper_limit': upper_limit
    }


"""
  Performs several operations and manipulations of geometries
"""

import numpy
from phydat import phycon
import automol.create.geom
from automol import util
from automol.geom import _base as geom_base
from automol.graph.geom import center_of_mass
from automol.graph.geom import translate as _translate
from automol.graph.geom import geometry_join as _geometry_join

AXIS_DCT = {'x': 0, 'y': 1, 'z': 2}

# General transformation functions
def transform(geo, func, idxs=None):
    """ Transform the coordinates of a geometry by a function.
        A set of `idxs` can be supplied to transform a subset of coordinates.

        :param geo: molecular geometry
        :type geo: automol geometry data structure
        :param func: transformation function
        :type func: function object
        :param idxs: indices representing the subset of atoms
        :type idxs: tuple(int)
    """

    idxs = list(range(geom_base.count(geo))) if idxs is None else idxs
    symbs = geom_base.symbols(geo)
    xyzs = geom_base.coordinates(geo)
    xyzs = [func(xyz) if idx in idxs else xyz for idx, xyz in enumerate(xyzs)]

    return automol.create.geom.from_data(symbs, xyzs)

def transform_by_matrix(geo, mat):
    """ Transform the coordinates of a molecular geometry by multiplying
        it by some input transfomration matrix.

        :param geo: molecular geometry
        :type geo: automol molecular geometry data structure
        :param mat: transformation matrix
        :type mat: tuple(tuple(float))
        :rtype: automol moleculer geometry data structure
    """

    symbs = geom_base.symbols(geo)
    xyzs = geom_base.coordinates(geo)
    xyzs = numpy.dot(xyzs, numpy.transpose(mat))

    return automol.create.geom.from_data(symbs, xyzs)

# transformations
def remove_coordinates(geo, idxs=()):
    """ Remove atoms by from the molecular geometry by their indices.

        :param geo: molecular geometry
        :type geo: automol molecular geometry data structure
        :param idxs: indices of atoms to remove
        :type idxs: tuple(int)
        :rtype: automol molecular geometry data structure
    """
    return tuple(row for i, row in enumerate(geo) if i not in idxs)

def join(geo1, geo2,
         dist_cutoff=3.0*phycon.ANG2BOHR, theta=0.0, phi=0.0):
    """ Join two molecular geometries together where the intermolecular
        separation and orientation can be specified.

        :param geo1: molecular geometry 1
        :type geo1: automol molecular geometry data structure
        :param geo2: molecular geometry 2
        :type geo2: automol molecular geometry data structure
        :param dist_cutoff: threshhold for center-of-mass distance
        :type: dist_cutoff: float
        :param theta: theta angle for intermolecular orientation
        :type theta: float
        :param phi: phi angle for intermolecular orientation
        :type phi: float
        :rtype: automol molecular geometry data structure
    """
    return _geometry_join(
        geo1, geo2, dist_cutoff=dist_cutoff, theta=theta, phi=phi)

def reorder_coordinates(geo, idx_dct):
    """ Reorder the atoms of a molecular geometry using
        the mapping of an input dictionary.

        :param geo: The geometry
        :param idx_dct: The new order of the atoms, by index
        :type idx_dct: dict
        :rtype: automol geometry data structure
    """

    symbs = geom_base.symbols(geo)
    xyzs = geom_base.coordinates(geo)

    idxs = [idx for idx, _ in sorted(idx_dct.items(), key=lambda x: x[1])]
    assert len(symbs) == len(xyzs) == len(idxs)

    symbs = [symbs[idx] for idx in idxs]
    xyzs = [xyzs[idx] for idx in idxs]

    return automol.create.geom.from_data(symbs, xyzs)

def move_atom(geo, idx1, idx2):
    """ Move an atom to a different position in the geometry

        :param geo: molecular geometry
        :type geo: automol molecular geometry data structure
        :param idx1: index of the atom to be moved
        :type idx1: int
        :param idx2: new position that the atom should be moved to
        :type idx2: int
        :returns: the transformed geometry
        :rtype: molecular geometry
    """
    symbs = list(geom_base.symbols(geo))
    xyzs = list(geom_base.coordinates(geo))
    symbs.insert(idx2, symbs.pop(idx1))
    xyzs.insert(idx2, xyzs.pop(idx1))
    return automol.create.geom.from_data(symbs, xyzs)

def swap_coordinates(geo, idx1, idx2):
    """ Swap the order of the coordinates of two atoms in a molecular geometry.

        :param geo: molecular geometry
        :type geo: automol molecular geometry data structure
        :param idx1: index for one atom to swap coordinates
        :type idx1: int
        :param idx2: index for one atom to swap coordinates
        :type idx2: int
        :rtype: molecular geometry
    """

    geo = [list(x) for x in geo]
    geo[idx1], geo[idx2] = geo[idx2], geo[idx1]
    geo_swp = tuple(tuple(x) for x in geo)

    return geo_swp

def insert(geo, symb, xyz, idx=None, angstrom=False):
    """ Insert an atom into a molecular geometry.

        :param geo: molecular geometry
        :type geo: automol molecular geometry data structure
        :param symb: symbol of atom to add
        :type symb: str
        :param xyz: xyz coordinates of atom to add
        :type xyz: tuple(float)
        :param idx: index of geometry to place atom
        :type idx: int
        :rtype: automol geometry date structure
    """

    return automol.convert.geom.insert(
        geo, symb, xyz, idx=idx, angstrom=angstrom)

def insert_dummies(geo, dummy_key_dct, dist=1., tol=5.):
    """ Insert dummy atoms over atoms in a geometry in a particular order.

        :param geo: the geometry
        :type geo: automol molecular geometry data structure
        :param dummy_key_dct: the linear atoms and the desired positions of the
            dummy atoms for each; linear atom indexing should follow what they
            *will* be after the dummy atoms are moved to the appropriate
            positions
        :param dummy_key_dct: dict
        :param dist: distance of dummy atom from the linear atom, in angstroms
        :type dist: float
        :param tol: the tolerance threshold for linearity, in degrees
        :type tol: float
        :returns: geometry with dummy atoms inserted, along with a dictionary
            mapping the linear atoms onto their associated dummy atoms
        :rtype: automol molecular geometry data structure
    """
    lin_keys, dum_keys = zip(
        *sorted(dummy_key_dct.items(), key=lambda x: x[1]))
    dum_keys = numpy.array(list(dum_keys))
    lin_idxs = [k-sum(k > dum_keys) for k in lin_keys]
    geo, orig_dummy_key_dct = insert_dummies_on_linear_atoms(
        geo, lin_idxs=lin_idxs, dist=dist, tol=tol)

    for lin_idx, lin_key in zip(lin_idxs, lin_keys):
        orig_idx = orig_dummy_key_dct[lin_idx]
        new_idx = dummy_key_dct[lin_key]
        geo = move_atom(geo, orig_idx, new_idx)

    return geo

def insert_dummies_on_linear_atoms(geo, lin_idxs=None, gra=None, dist=1.,
                                   tol=5.):
    """ Insert dummy atoms over linear atoms in the geometry.

        :param geo: the geometry
        :type geo: automol molecular geometry data structure
        :param lin_idxs: the indices of the linear atoms; if None, indices are
            automatically determined from the geometry based on the graph
        :type lin_idxs: tuple(int)
        :param gra: the graph describing connectivity; if None, a connectivity
            graph will be generated using default distance thresholds
        :type gra: automol molecular graph data structure
        :param dist: distance of dummy atom from the linear atom, in angstroms
        :type dist: float
        :param tol: the tolerance threshold for linearity, in degrees
        :type tol: float
        :returns: geometry with dummy atoms inserted, along with a dictionary
            mapping the linear atoms onto their associated dummy atoms
        :rtype: automol molecular geometry data structure
    """
    return automol.convert.geom.insert_dummies_on_linear_atoms(
        geo, lin_idxs=lin_idxs, gra=gra, dist=dist, tol=tol)

def displace(geo, xyzs):
    """ Displace the coordinates of a geometry along a vector.

        :param geo: molecular geometry
        :type geo: automol molecular geometry data structure
        :param xyzs: vector to displace along
        :type xyzs: tuple(float)
        :rtype: automol molecular geometry data structure
    """

    symbs = geom_base.symbols(geo)
    orig_xyzs = geom_base.coordinates(geo)
    xyzs = numpy.add(orig_xyzs, xyzs)

    return automol.create.geom.from_data(symbs, xyzs)

# redundant with above
def translate(geo, xyz):
    """ Translate the coordinates of a molecular geometry along
        a three-dimensiona vector.

        :param geo: molecular geometry
        :type geo: automol molecular geometry data structure
        :param xyz: vector to translate along
        :type xyz: tuple(float)
        :rtype: automol molecular geometry data structure
    """
    return _translate(geo, xyz)

def perturb(geo, atm_idx, pert_xyz):
    """ Perturb the position of one atom by
        changing the value of an xyz coord by some amount.
    """

    # Get the xyz coordinates of the atom to perturb
    atm_coords = list(geom_base.coordinates(geo)[atm_idx])

    # Get the perturbed set of atomic coordinates
    for idx, val in enumerate(pert_xyz):
        atm_coords[idx] += val
    pert_dct = {atm_idx: atm_coords}

    # Perturb the coordinates of the atom
    pert_geo = geom_base.set_coordinates(geo, pert_dct)

    return pert_geo

def mass_centered(geo):
    """ Generate a new geometry where the coordinates of the input geometry
        have been translated to the center-of-mass.

        :param geo: molecular geometry
        :type geo: automol geometry data structure
        :rtype: tuple(float)
    """
    return translate(geo, numpy.negative(center_of_mass(geo)))

def rotate(geo, axis, angle, orig_xyz=None, idxs=None):
    """ Rotate the coordinates of a molecular geometry about
        an axis by a specified angle. A set of `idxs` can be supplied
        to transform a subset of coordinates.

        :param geo: molecular geometry
        :type geo: automol geometry data structure
        :param axis: axis to rotate about
        :type axis: tuple(float)
        :param angle: angle of rotation
        :type angle: float
        :param orig_xyz: xyz coordinates of the origin
        :type orig_xyz: tuple(float)
        :param idxs: indices of atoms whose coordinates are to be rotated
        :type idxs: tuple(int)
    """

    func = util.vec.rotater(axis, angle, orig_xyz=orig_xyz)

    return transform(geo, func, idxs=idxs)

def euler_rotate(geo, theta, phi, psi):
    """ Rotate the coordinates of a molecular geometry about
        the three Euler angles.

        :param geo: molecular geometry
        :type geo: automol geometry data structure
        :param theta: angle to rotate about z-axis
        :type theta: float
        :param phi: angle to rotate about x'-axis
        :type phi: float
        :param psi: angle to rotate about z'-axis
        :type psi: float
        :rtype: automol geometry data structure
    """

    mat = util.mat.euler_rotation_matrix(theta, phi, psi)

    return transform_by_matrix(geo, mat)

def shift_atom_position(geo, idx1, idx2):
    """ Move the atom at position idx1 to idx2, shifting all other atoms.

        :param geo: molecular geometry
        :type geo: automol geometry data structure
        :param idx1: index of atom 1 in the pair to be measured
        :type idx1: int
        :param idx2: index of atom 2 in the pair to be measured
        :type idx2: int
        :rtype: automol geometry data structure
    """

    # Get the coordinates at idx1 that are to be moved
    geo = [list(x) for x in geo]
    moving_coords = geo[idx1]

    # move the coordinates to idx2
    geo.remove(moving_coords)
    geo.insert(idx2, moving_coords)
    geo_move = tuple(tuple(x) for x in geo)

    return geo_move

def reflect_coordinates(geo, idxs, axes):
    """ Reflect a specified set of coordinates of a molecular geometry
        about some each of the requested axes.

        A set of `idxs` can be supplied to transform a subset of coordinates.

        :param geo: molecular geometry
        :type geo: automol geometry data structure
        :param idxs: indices of atoms whose coordinates are to be reflected
        :type idxs: tuple(int)
        :param axes: axes to reflect about
        :type axes: tuple(str)
        :rtype: automol geometry data structure
    """

    # check input
    assert all(idx < len(geo) for idx in idxs)
    assert all(axis in ('x', 'y', 'z') for axis in axes)

    # get coords
    coords = geom_base.coordinates(geo)

    # convert x,y,z to nums
    axes = [AXIS_DCT[axis] for axis in axes]

    # build set atom dct with relected coords
    reflect_dct = {}
    for idx in idxs:
        coord_lst = list(coords[idx])
        for axis in axes:
            coord_lst[axis] *= -1.0
        reflect_dct[idx] = coord_lst

    # Reflect coords with dct
    geo_reflected = geom_base.set_coordinates(geo, reflect_dct)

    return geo_reflected


import os
import os.path as osp
import json
from tqdm import tqdm
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Data, InMemoryDataset
from itertools import repeat
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit import DataStructs
from chem import *

class MultiGraphData(Data):
    def __init__(self, x=None, edge_index=None, edge_attr=None,
                 fg_x=None, fg_edge_index=None, atom2fg_index=None,
                 **kwargs):
        super(MultiGraphData, self).__init__(x=x, edge_index=edge_index, edge_attr=edge_attr)
        self.fg_x = fg_x
        self.fg_edge_index = fg_edge_index
        self.num_fgs = fg_x.size(0) if fg_x is not None else None
        self.atom2fg_index = atom2fg_index
        for key, value in kwargs.items():
            setattr(self, key, value)

    def __inc__(self, key, value):
        r"""Returns the incremental count to cumulatively increase the value
        of the next attribute of :obj:`key` when creating batches.

        .. note::

            This method is for internal use only, and should only be overridden
            if the batch concatenation process is corrupted for a specific data
            attribute.
        """
        # Only `*index*` and `*face*` attributes should be cumulatively summed
        # up when creating batches.
        if key == 'fg_edge_index':
            return self.fg_x.size(0)
        elif key == 'atom2fg_index':
            return torch.tensor([[self.num_nodes], [self.fg_x.size(0)]])
        else:
            return super(MultiGraphData, self).__inc__(key, value)

class PretrainDataset(InMemoryDataset):
    def __init__(self, root='data/ZINC15',
                 mol_filename='zinc15_250k.txt',
                 fg_corpus_filename='fg_corpus.txt',
                 mol2fgs_filename='mol2fgs_list.json',
                 ):
        self.mol_fn = mol_filename
        self.corpus_fn = fg_corpus_filename
        self.mol2fgs_fn = mol2fgs_filename
        super().__init__(root=root)

        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_dir(self):
        return self.root

    @property
    def raw_file_names(self):
        return [self.mol_fn, self.corpus_fn, self.mol2fgs_fn]

    @property
    def processed_file_names(self):
        return osp.splitext(self.raw_file_names[0])[0] + '.pt'

    def get(self, idx):
        data = self.data.__class__()
        if hasattr(self.data, '__num_nodes__'):
            data.num_nodes = self.data.__num_nodes__[idx]

        for key in self.data.keys:
            item, slices = self.data[key], self.slices[key]
            start, end = slices[idx].item(), slices[idx + 1].item()
            if torch.is_tensor(item):
                s = list(repeat(slice(None), item.dim()))
                s[self.data.__cat_dim__(key, item)] = slice(start, end)
            elif start + 1 == end:
                s = slices[start]
            else:
                s = slice(start, end)
            data[key] = item[s]

        return data

    def process(self):
        with open(self.raw_paths[0], 'r') as f:
            smiles_list = f.read().splitlines()
        with open(self.raw_paths[1], 'r') as f:
            fg_corpus = f.read().splitlines()
        with open(self.raw_paths[2], 'r') as f:
            mol2fgs = json.load(f)
        print(f"# mol: {len(smiles_list)}   # corpus: {len(fg_corpus)}")

        data_list = []
        for smiles, fgs in tqdm(zip(smiles_list, mol2fgs)):
            mol = Chem.MolFromSmiles(smiles)
            atom_features, bond_list, bond_features, fg_features, fg_edge_list, fg_edge_features, atom2fg_list = mol_to_graphs(mol)
            fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=1024)
            fpvec = np.zeros(0)
            DataStructs.ConvertToNumpyArray(fp, fpvec)
            fgvec = np.zeros(len(fg_corpus))
            idx = []
            for fg in fgs:
                try:
                    idx.append(fg_corpus.index(fg))
                except:
                    pass
            fgvec[idx] = 1
            data = MultiGraphData(x=torch.Tensor(atom_features),
                                  edge_index=torch.LongTensor(bond_list).reshape(-1, 2).transpose(1, 0),
                                  edge_attr=torch.Tensor(bond_features).reshape(-1, BOND_DIM),
                                  fg_x=torch.Tensor(fg_features),
                                  fg_edge_index=torch.LongTensor(fg_edge_list).reshape(-1, 2).transpose(1, 0),
                                  fg_edge_attr=torch.Tensor(fg_edge_features).reshape(-1, FG_EDGE_DIM),
                                  atom2fg_index=torch.LongTensor(atom2fg_list).reshape(-1, 2).transpose(1, 0),
                                  fp=torch.Tensor(fpvec).reshape(1, -1),
                                  fg=torch.Tensor(fgvec).reshape(1, -1))
            data_list.append(data)

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class MoleculeNetDataset():
    def __init__(self, dataset):
        self.dataset = dataset

    def __getitem__(self, index):
        atom_features, bond_list, bond_features, fg_features, fg_edge_list, fg_edge_features, atom2fg_list, y, w = self.dataset[index]
        data = MultiGraphData(x=torch.Tensor(atom_features),
                              edge_index=torch.LongTensor(bond_list).reshape(-1, 2).transpose(1, 0),
                              edge_attr=torch.Tensor(bond_features).reshape(-1, BOND_DIM),
                              fg_x=torch.Tensor(fg_features),
                              fg_edge_index=torch.LongTensor(fg_edge_list).reshape(-1, 2).transpose(1, 0),
                              fg_edge_attr=torch.Tensor(fg_edge_features).reshape(-1, FG_EDGE_DIM),
                              atom2fg_index=torch.LongTensor(atom2fg_list).reshape(-1, 2).transpose(1, 0),
                              y=torch.Tensor(y),
                              w=torch.Tensor(w))
        return data

    def __len__(self):
        return len(self.dataset)

class DDIDataset(InMemoryDataset):
    def __init__(self, root='data/DDI/ZhangDDI',
                 drug_filename='drug_list_zhang.csv',
                 ddi_filename='ZhangDDI_train.csv'):
        self.drug_fn = drug_filename
        self.ddi_fn = ddi_filename
        super().__init__(root=root)

        self.drugs = torch.load(self.processed_paths[0])
        df = pd.read_csv(os.path.join(self.root, self.ddi_fn), usecols=['smiles_1', 'smiles_2', 'label'])
        self.ddi = df.values

    @property
    def raw_dir(self):
        return self.root

    @property
    def raw_file_names(self):
        return self.drug_fn

    @property
    def processed_file_names(self):
        return osp.splitext(self.raw_file_names)[0] + '.pt'

    def __getitem__(self, idx):
        id1, id2, label = self.ddi[idx]
        return self.drugs[id1], self.drugs[id2], torch.Tensor([float(label)])

    def __len__(self):
        return len(self.ddi)

    def process(self):
        df = pd.read_csv(self.raw_paths[0], usecols=['drugbank_id', 'smiles'])
        print(f"# drugs: {len(df)}")

        data_dict = {}
        for _, drug in tqdm(df.iterrows()):
            id, smiles = drug['drugbank_id'], drug['smiles']
            mol = Chem.MolFromSmiles(smiles)
            atom_features, bond_list, bond_features, fg_features, fg_edge_list, fg_edge_features, atom2fg_list = mol_to_graphs(mol)
            if fg_features == []:  # C
                print(f"{smiles} cannot be converted to FG graph")
                continue
            data = MultiGraphData(x=torch.Tensor(atom_features),
                                  edge_index=torch.LongTensor(bond_list).reshape(-1, 2).transpose(1, 0),
                                  edge_attr=torch.Tensor(bond_features).reshape(-1, BOND_DIM),
                                  fg_x=torch.Tensor(fg_features),
                                  fg_edge_index=torch.LongTensor(fg_edge_list).reshape(-1, 2).transpose(1, 0),
                                  fg_edge_attr=torch.Tensor(fg_edge_features).reshape(-1, FG_EDGE_DIM),
                                  atom2fg_index=torch.LongTensor(atom2fg_list).reshape(-1, 2).transpose(1, 0))
            data_dict[smiles] = data

        torch.save(data_dict, self.processed_paths[0])

if __name__ == '__main__':
    dataset = PretrainDataset(mol_filename='zinc15_250k.txt',
                              fg_corpus_filename='fg_corpus.txt',
                              mol2fgs_filename='mol2fgs_list.json')
    data = dataset[0]