"""The module contains everything to handle cross section interfaces."""

from abc import ABCMeta, abstractmethod
from os.path import isfile, join

import numpy as np

import prince_cr.decays as decs
from prince_cr.data import spec_data
from prince_cr.util import (
    get_2Dinterp_object, get_interp_object, info, bin_widths, get_AZN)
import prince_cr.config as config

class CrossSectionBase(object, metaclass=ABCMeta):
    """Base class for cross section interfaces to tabulated models.

    The class is abstract and it is not inteded to be instantiated.

    Child Classes either define the tables:
        self._egrid_tab
        self._nonel_tab
        self._incl_tab
        self._incl_diff

    Or directly reimplememnt the functions
        nonel(self, mother, daughter)
        incl(self, mother, daughter)
        incl_diff(self, mother, daughter)

    The flag self.supports_redistributions = True/False should be set
    To tell the class to include/ignore incl_diff
    """
    def __init__(self):
        # Tuple, defining min and max energy cuts on the grid
        self._range = None
        # Energy grid, as defined in files
        self._egrid_tab = None
        # Dictionary of nonel. cross sections on egrid, indexed by (mother)
        self._nonel_tab = {}
        # Dictionary of incl. cross sections on egrid, indexed by (mother, daughter)
        self._incl_tab = {}
        # Dictionary of incl. diff. cross sections on egrid, indexed by (mother, daughter)
        self._incl_diff_tab = {}
        # List of available mothers for nonel cross sections
        self.nonel_idcs = []
        # List of available (mothers,daughter) reactions in incl. cross sections
        self.incl_idcs = []
        # List of available (mothers,daughter) reactions in incl. diff. cross sections
        self.incl_diff_idcs = []
        # Common grid in x (the redistribution variable)
        self.xbins = None

        # Flag, which tells if the model supports secondary redistributions
        if not hasattr(self, 'supports_redistributions'):
            self.supports_redistributions = None  # JH: to differ from explicitly set False
        # List of all known particles (after optimization)
        self.known_species = []
        # List of all boost conserving inclusive channels (after optimization)
        self.known_bc_channels = []
        # List of all differential inclusive channels (after optimization)
        self.known_diff_channels = []
        # Dictionary of (mother, daughter) reactions for each mother
        self.reactions = {}

        # Class name of the model
        self.mname = self.__class__.__name__

    def set_range(self, e_min=None, e_max=None):
        """Set energy range within which to return tabulated data.

        Args:
            e_min (float): minimal energy in GeV
            e_max (float): maximal energy in GeV
        """
        if e_min is None:
            e_min = np.min(self._egrid_tab)
        if e_max is None:
            e_max = np.max(self._egrid_tab)

        info(5, "Setting range to {0:3.2e} - {1:3.2e}".format(e_min, e_max))
        self._range = np.where((self._egrid_tab >= e_min)
                               & (self._egrid_tab <= e_max))[0]
        info(
            2, "Range set to {0:3.2e} - {1:3.2e}".format(
                np.min(self._egrid_tab[self._range]),
                np.max(self._egrid_tab[self._range])))

    @property
    def egrid(self):
        """Returns energy grid of the tabulated data in selected range.

        Returns:
            (numpy.array): Energy grid in GeV
        """

        return self._egrid_tab[self._range]

    @property
    def xcenters(self):
        """Returns centers of the grid in x.

        Returns:
            (numpy.array): x grid
        """

        return 0.5 * (self.xbins[1:] + self.xbins[:-1])

    @property
    def xwidths(self):
        """Returns bin widths of the grid in x.

        Returns:
            (numpy.array): x widths
        """

        return self.xbins[1:] - self.xbins[:-1]

    @property
    def resp(self):
        """Return ResponseFunction corresponding to this cross section
        Will only create the Response function once.
        """
        if not hasattr(self, '_resp'):
            info(2, 'First Call, creating instance of ResponseFunction now')
            from .response import ResponseFunction
            self._resp = ResponseFunction(self)
        return self._resp

    def is_differential(self, mother, daughter):
        """Returns true if the model supports redistributions and requested
        mother/daughter combination should return non-zero redistribution matrices.

        Args:
            mother (bool): Neucosma ID of mother particle
            daughter (bool): Neucosma ID of daughter particle

        Returns:
            (bool): ``True`` if the model has this particular redistribution function
        """
        # info(10, mother, daughter, " asking for redist")
        # if not self.supports_redistributions:
        #     info(10, mother, daughter, " model doesn't support redist")
        #     return False
        if (daughter <= config.redist_threshold_ID
                or (mother, daughter) in self.incl_diff_idcs):
            info(60, 'Daughter requires redistribution.', mother, daughter)
            return True
        info(60, 'Daughter conserves boost.', mother, daughter)
        return False

    def _update_indices(self):
        """Updates the list of indices according to entries in the
        _tab variables"""

        self.nonel_idcs = sorted(self._nonel_tab.keys())
        self.incl_idcs = sorted(self._incl_tab.keys())
        self.incl_diff_idcs = sorted(self._incl_diff_tab.keys())

    def generate_incl_channels(self, mo_indices):
        """Generates indices for all allowed channels given mo_indices
            Note: By default this returns an empty list,
                  meant to be overwritten in cases where
                  the child class needs to dynamically generate indices

        Args:
            mo_indices (list of ints): list of indices for mother nuclei

        Returns:
           Returns:
            list of tuples: list of allowed channels given as (mo_idx, da_idx)
        """
        incl_channels = []

        return incl_channels

    def _optimize_and_generate_index(self):
        """Construct a list of mothers and (mother, daughter) indices.

        Args:
            just_reactions (bool): If True then fill just the reactions index.
        """

        # Integrate out short lived processes and leave only stable particles
        # in the databases
        self._reduce_channels()

        # Go through all three cross section categories
        # index contents in the ..known..variable
        self.reactions = {}

        self._update_indices()

        for mo, da in self.incl_idcs:
            if da >= 100 and get_AZN(da)[0] > get_AZN(mo)[0]:
                raise Exception(
                    'Daughter {0} heavier than mother {1}. Physics??'.format(
                        da, mo))

            if mo not in self.reactions:
                self.reactions[mo] = []
                self.known_species.append(mo)

            if (mo, da) not in self.reactions[mo]:
                # Make sure it's a unique list
                self.reactions[mo].append((mo, da))
            if self.is_differential(mo, da):
                # Move the distributions which are expected to be differential
                # to _incl_diff_tab
                self._incl_diff_tab[(mo, da)] = self._arange_on_xgrid(
                    self._incl_tab.pop((mo, da)))
                info(10, "Channel {0} -> {1} forced to be differential.")
            else:
                self.known_bc_channels.append((mo, da))
                self.known_species.append(da)

        for mo, da in list(self._incl_diff_tab.keys()):
            if da >= 100 and get_AZN(da)[0] > get_AZN(mo)[0]:
                raise Exception(
                    'Daughter {0} heavier than mother {1}. Physics??'.format(
                        da, mo))

            if mo not in self.reactions:
                self.reactions[mo] = []
                self.known_species.append(mo)

            if (mo, da) not in self.reactions[mo]:
                # Make sure it's a unique list to avoid unnecessary loops
                self.reactions[mo].append((mo, da))
                self.known_diff_channels.append((mo, da))
                self.known_species.append(da)

        # Remove duplicates
        self.known_species = sorted(list(set(self.known_species)))
        self.known_bc_channels = sorted(list(set(self.known_bc_channels)))
        self.known_diff_channels = sorted(list(set(self.known_diff_channels)))

        for sp in self.known_species:
            if sp >= 100 and (sp, sp) not in self.known_diff_channels:
                self.known_bc_channels.append((mo, mo))
            if (mo, mo) not in self.reactions[mo]:
                self.reactions[mo].append((mo, mo))

        # Make sure the indices are up to date
        self._update_indices()

        # Count numbers of channels for statistics
        # Count number of incl channels for activated nuclear species
        # n_incl = np.sum([
        #     len(self.reactions[mother])
        #     for mother in self.spec_man.known_species if mother >= 100
        # ])

    def _reduce_channels(self):
        """Follows decay chains until all inclusive reactions point to
        stable final state particles.

        The "tau_dec_threshold" parameter in the config controls the
        definition of stable. Unstable nuclei for which no decay channels
        are known, will be forced to beta-decay until they reach a stable
        element.
        """
        from prince_cr.util import AdditiveDictionary
        # TODO: check routine, how to avoid empty channels and
        # mothers with zero nonel cross sections

        # The new dictionary that will replace _incl_tab
        new_incl_tab = AdditiveDictionary()
        new_dec_diff_tab = AdditiveDictionary()

        threshold = config.tau_dec_threshold

        # How to indent debug printout for recursion
        dbg_indent = lambda lev: 4 * lev * "-" + ">" if lev else ""

        info(2, "Integrating out species with lifetime smaller than",
             threshold)
        info(3, (
            "Before optimization, the number of known primaries is {0} with " +
            "in total {1} inclusive channels").format(len(self._nonel_tab),
                                                      len(self._incl_tab)))

        if self.xbins is None:
            info(
                4,
                'Model does not provide a native xbins. Assuming JH special sophia',
                'binning.')
            from .photo_meson import SophiaSuperposition
            self.xbins = SophiaSuperposition().xbins

        bc = self.xcenters
        bw = bin_widths(self.xbins)
        # The x_mu/x_pi grid
        # dec_grid = np.fromfunction(
        #     lambda j, i: 10**(np.log10(bc[1] / bc[0]) * (j - i)), (len(bc),
        #                                                            len(bc)))

        # dec_grid = np.outer(bc, 1 / bc)

        dec_bins = np.outer(self.xbins, 1 / bc)
        dec_bins_lower = dec_bins[:-1]
        dec_bins_upper = dec_bins[1:]

        # dec_grid[dec_grid > 1.] *= 0.
        # The differential element dx_mu/x_pi
        int_scale = np.tile(bw / bc, (len(bc), 1))

        from functools import lru_cache
        @lru_cache(maxsize=512, typed=False)
        def decay_cached(mother,daughter):
            dec_dist = int_scale * decs.get_decay_matrix_bin_average(
                mother, daughter, dec_bins_lower, dec_bins_upper)

            return dec_dist

        def convolve_with_decay_distribution(diff_dist, mother, daughter,
                                             branching_ratio):
            r"""Computes the prompt decay xdist by convolving the x distribution
            of the unstable particle with the decay product distribution.

            :math:`\frac{{\rm d}N^{A\gamma \to \mu}}{{\rm d}x_j} =
            \sum_{i=0}^{N_x}~\Delta x_i
            \frac{{\rm d}N^{A\gamma \to \pi}}{{\rm d} x_i}~
            \frac{{\rm d}N^{\pi \to \mu}}{{\rm d} x_j}`
            """
            # dec_dist = int_scale * decs.get_decay_matrix(
            #     mother, daughter, dec_grid)
            dec_dist = decay_cached(mother,daughter)

            info(20, 'convolving with decay dist', mother, daughter)
            # Handle the case where table entry is (energy_grid, matrix)
            if not isinstance(diff_dist, tuple):
                return branching_ratio * dec_dist.dot(diff_dist)
            else:
                return diff_dist[0], branching_ratio * dec_dist.dot(
                    diff_dist[1])

        def follow_chain(first_mo, da, csection, reclev):
            """Recursive function to follow decay chains until all
            final state particles are stable.

            The result is saved in two dictionaries; one for the boost
            conserving inclusive channels and the other one collects
            channels with meson or lepton decay products, which will
            need special care due to energy redistributions of these
            secondaries.
            """

            info(10, dbg_indent(reclev), 'Entering with', first_mo, da)

            if da not in spec_data:
                info(
                    3, dbg_indent(reclev),
                    'daughter {0} unknown, forcing beta decay. Not Implemented yet!!'
                    .format(da))
                return

            # Daughter is stable. Add it to the new dictionary and terminate
            # recursion
            if spec_data[da]["lifetime"] >= threshold:
                if self.is_differential(None, da):
                    # If the daughter is a meson or lepton, use the dictionary for
                    # differential channels
                    info(
                        20, dbg_indent(reclev),
                        'daughter {0} stable and differential. Adding to ({1}, {2})'
                        .format(da, first_mo, da))
                    new_dec_diff_tab[(first_mo, da)] = csection
                else:
                    info(
                        20, dbg_indent(reclev),
                        'daughter {0} stable. Adding to ({1}, {2})'.format(
                            da, first_mo, da))
                    new_incl_tab[(first_mo, da)] = csection
                return

            # ..otherwise follow decay products of this daughter, tracking the
            # original mother particle (first_mo). The cross section (csection) is
            # reduced by the branching ratio (br) of this particular channel
            for br, daughters in spec_data[da]["branchings"]:
                info(10, dbg_indent(reclev),
                     ("{3} -> {0:4d} -> {2:4.2f}: {1}").format(
                         da, ", ".join(map(str, daughters)), br, first_mo))

                for chained_daughter in daughters:
                    # Follow each secondary and increment the recursion level by one
                    if self.is_differential(None, chained_daughter):
                        info(10, 'daughter', chained_daughter, 'of', da,
                             'is differential')
                        follow_chain(
                            first_mo, chained_daughter,
                            convolve_with_decay_distribution(
                                self._arange_on_xgrid(csection), da,
                                chained_daughter, br), reclev + 1)
                    else:
                        follow_chain(first_mo, chained_daughter, br * csection,
                                     reclev + 1)

        # Remove all unstable particles from the dictionaries
        for mother in sorted(self._nonel_tab.keys()):
            if mother not in spec_data or spec_data[mother][
                    "lifetime"] < threshold:
                info(
                    20,
                    "Primary species {0} does not fulfill stability criteria.".
                    format(mother))
                _ = self._nonel_tab.pop(mother)
        # Only stable (interacting) mother particles are left
        self._update_indices()

        for (mother, daughter) in self.incl_idcs:

            if mother not in self.nonel_idcs:
                info(
                    30, "Removing {0}/{1} from incl, since mother not stable ".
                    format(mother, daughter))
                _ = self._incl_tab.pop((mother, daughter))

            elif self.is_differential(mother, daughter):
                # Move the distributions which are expected to be differential
                # to _incl_diff_tab
                self._incl_diff_tab[(mother,
                                     daughter)] = self._arange_on_xgrid(
                                         self._incl_tab.pop(
                                             (mother, daughter)))

        self._update_indices()

        for (mother, daughter) in self.incl_diff_idcs:

            if mother not in self.nonel_idcs:
                info(
                    30,
                    "Removing {0}/{1} from diff incl, since mother not stable "
                    .format(mother, daughter))
                _ = self._incl_diff_tab.pop((mother, daughter))

        self._update_indices()

        # Launch the reduction for each inclusive channel
        for (mo, da), value in list(self._incl_tab.items()):
            #print mo, da, value
            #print '---'*30
            follow_chain(mo, da, value, 0)

        for (mo, da), value in list(self._incl_diff_tab.items()):
            #print mo, da, value
            #print '---'*30
            follow_chain(mo, da, value, 0)

        # Overwrite the old incl dictionary
        self._incl_tab = dict(new_incl_tab)
        # Overwrite the old incl_diff dictionary
        self._incl_diff_tab = dict(new_dec_diff_tab)
        # Reduce also the incl_diff_tab by removing the unknown mothers. At this stage
        # of the code, the particles with redistributions are
        info(
            3,
            ("After optimization, the number of known primaries is {0} with " +
             "in total {1} inclusive channels").format(
                 len(self._nonel_tab),
                 len(self._incl_tab) + len(self._incl_diff_tab)))
        info(2, f'Cache used for decays, {decay_cached.cache_info()}') # pylint:disable=no-value-for-parameter

    def nonel_scale(self, mother, scale='A'):
        """Returns the nonel cross section scaled by `scale`.

        Convenience funtion for plotting, where it is important to
        compare the cross section per nucleon.

        Args:
            mother (int): Mother nucleus(on)
            scale (float): If `A` then nonel/A is returned, otherwise
                           scale can be any float.

        Returns:
            (numpy.array, numpy.array): Tuple of Energy grid in GeV,
                                        scale * inclusive cross section
                                        in :math:`cm^{-2}`
        """

        egr, csection = self.nonel(mother)

        if scale == 'A':
            scale = 1. / get_AZN(mother)[0]

        return egr, scale * csection

    def incl_scale(self, mother, daughter, scale='A'):
        """Same as :func:`~cross_sections.CrossSectionBase.nonel_scale`,
        just for inclusive cross sections.
        """

        egr, csection = self.incl(mother, daughter)

        if scale == 'A':
            scale = 1. / get_AZN(mother)[0]

        return egr, scale * csection

    def nonel(self, mother):
        """Returns non-elastic cross section.

        Absorption cross section of `mother`, which is
        the total minus elastic, or in other words, the inelastic
        cross section.

        Args:
            mother (int): Mother nucleus(on)

        Returns:
            (numpy.array, numpy.array): Tuple of Energy grid in GeV, inclusive cross
                                        section in :math:`cm^{-2}`
        """

        if mother not in self._nonel_tab:
            raise Exception('Mother {0} unknown.'.format(mother))

        if isinstance(self._nonel_tab[mother], tuple):
            return self._nonel_tab[mother]
        else:
            return self.egrid, self._nonel_tab[mother][self._range]

    def incl(self, mother, daughter):
        """Returns inclusive cross section.

        Inclusive cross section for daughter in photo-nuclear
        interactions of `mother`.

        Args:
            mother (int): Mother nucleus(on)
            daughter (int): Daughter nucleus(on)

        Returns:
            (numpy.array): Inclusive cross section in :math:`cm^{-2}`
                           on self._egrid_tab
        """

        from scipy.integrate import trapz

        if (mother, daughter) in self._incl_diff_tab:
            # Return the integral of the differential for the inclusive
            egr_incl, cs_diff = self.incl_diff(mother, daughter)
            # diff_mat = diff_mat.transpose()
            cs_incl = trapz(cs_diff,
                            x=self.xcenters,
                            dx=bin_widths(self.xbins),
                            axis=0)

            if isinstance(self._incl_diff_tab[(mother, daughter)], tuple):
                return egr_incl, cs_incl

            return self.egrid, cs_incl[self._range]

        elif (mother, daughter) not in self._incl_tab:
            raise Exception(
                self.__class__.__name__ + '::'
                '({0},{1}) combination not in inclusive cross sections'.format(
                    mother, daughter))

        # If _nonel_tab contains tuples of (egrid, cs) return tuple
        # otherwise return (egrid, cs) in range defined by self.range

        if isinstance(self._incl_tab[(mother, daughter)], tuple):
            return self._incl_tab[(mother, daughter)]
        return self.egrid, self._incl_tab[(mother, daughter)][self._range]

    def incl_diff(self, mother, daughter):
        """Returns inclusive cross section.

        Inclusive differential cross section for daughter in photo-nuclear
        interactions of `mother`. Only defined, if the daughter is distributed
        in :math:`x = E_{da} / E_{mo}`

        Args:
            mother (int): Mother nucleus(on)
            daughter (int): Daughter nucleus(on)

        Returns:
            (numpy.array): Inclusive cross section in :math:`cm^{-2}`
                           on self._egrid_tab
        """

        if (mother, daughter) not in self._incl_diff_tab:
            raise Exception(
                self.__class__.__name__ +
                '({0},{1}) combination not in inclusive differential cross sections'
                .format(mother, daughter))

        # If _nonel_tab contains tuples of (egrid, cs) return tuple
        # otherwise return (egrid, cs) in range defined by self.range

        if isinstance(self._incl_diff_tab[(mother, daughter)], tuple):
            return self._incl_diff_tab[(mother, daughter)]
        return self.egrid, self._incl_diff_tab[(mother,
                                                daughter)][:, self._range]

    def _arange_on_xgrid(self, incl_cs):
        """Returns the inclusive cross section on an xgrid at x=1."""

        egr, cs = None, None

        if isinstance(incl_cs, tuple):
            egr, cs = incl_cs
        else:
            cs = incl_cs

        nxbins = len(self.xbins) - 1
        if len(cs.shape) > 1 and cs.shape[0] != nxbins:
            raise Exception(
                'One dimensional cross section expected, instead got',
                cs.shape, '\n', cs)
        elif len(cs.shape) == 2 and cs.shape[0] == nxbins:
            info(20, 'Supplied 2D distribution seems to be distributed in x.')
            if isinstance(incl_cs, tuple):
                return egr, cs
            return cs

        csec = np.zeros((nxbins, cs.shape[0]))
        # NOTE: The factor 2 in the following line is a workarround to account for the latter linear interpolation
        #       This is needed because linear spline integral will result in a trapz,
        #       which has half the area of the actual first bin
        corr_factor = 2 * self.xwidths[-1] / (self.xcenters[-1] -
                                              self.xcenters[-2])
        csec[-1, :] = cs / self.xwidths[-1] * corr_factor
        info(
            4,
            'Warning! Workaround to account for linear interpolation in x, factor 2 added!'
        )
        if isinstance(incl_cs, tuple):
            return egr, csec
        return csec

    def multiplicities(self, mother, daughter):
        '''Return the multiplicities from either the inclusive channels, or the
        differential ones integrated by x, as a function of Energy.
        '''
        egrid_incl, cs_incl = self.incl(mother, daughter)
        egrid_nonel, cs_nonel = self.nonel(mother)

        if egrid_incl.shape != egrid_nonel.shape:
            raise Exception('Problem with different grid shapes')

        multiplicities = cs_incl / np.where(cs_nonel == 0, np.inf, cs_nonel)

        return egrid_nonel, multiplicities

if __name__ == "__main__":
    pass


"""
Module for handling atomic elements.
"""
from numpy import vectorize

class Element:
    """
    Class for storing and accessing atomic elements.

    Args:
        input_value: The atomic number, symbol, or name of the element
    """

    def __init__(self, input_value):
        self.input = input_value

        # list with atomic number z, short name, full name, valence,
        # valence electrons, covalent radius, vdW radius, metallic radius
        self.elements_list = [
            (1, "H", "Hydrogen", 1.0, 1, 0.31, 1.20, None),
            (2, "He", "Helium", 0.5, 2, 0.28, 1.40, None),
            (3, "Li", "Lithium", 1.0, 1, 1.28, 1.82, 1.52),
            (4, "Be", "Beryllium", 2.0, 2, 0.96, 1.53, 1.12),
            (5, "B", "Boron", 3.0, 3, 0.84, 1.92, None),
            (6, "C", "Carbon", 4.0, 4, 0.78, 1.70, None),
            (7, "N", "Nitrogen", 3.0, 5, 0.78, 1.55, None),
            (8, "O", "Oxygen", 2.0, 6, 0.70, 1.52, None),
            (9, "F", "Fluorine", 1.0, 7, 0.57, 1.47, None),
            (10, "Ne", "Neon", 0.5, 8, 0.58, 1.54, None),
            (11, "Na", "Sodium", 1.0, 1, 1.66, 2.27, 1.86),
            (12, "Mg", "Magnesium", 2.0, 2, 1.41, 1.73, 1.60),
            (13, "Al", "Aluminium", 3.0, 3, 1.21, 1.84, 1.43),
            (14, "Si", "Silicon", 4.0, 4, 1.11, 2.10, None),
            (15, "P", "Phosphorus", 3.0, 5, 1.07, 1.80, None),
            (16, "S", "Sulfur", 2.0, 6, 1.25, 1.80, None),
            (17, "Cl", "Chlorine", 1.0, 7, 1.02, 1.75, None),
            (18, "Ar", "Argon", 0.5, 8, 1.06, 1.88, None),
            (19, "K", "Potassium", 1.0, 1, 2.03, 2.75, 2.27),
            (20, "Ca", "Calcium", 2.0, 2, 1.76, 2.31, 1.97),
            (21, "Sc", "Scandium", 3.0, 3, 1.70, 2.11, 1.62),
            (22, "Ti", "Titanium", 4.0, 4, 1.60, 2.00, 1.47),
            (23, "V", "Vanadium", 4.0, 5, 1.53, 2.00, 1.34),
            (24, "Cr", "Chromium", 3.0, 6, 1.39, 2.00, 1.28),
            (25, "Mn", "Manganese", 4.0, 5, 1.39, 2.00, 1.27),
            (26, "Fe", "Iron", 3.0, 3, 1.32, 2.00, 1.26),
            (27, "Co", "Cobalt", 3.0, 3, 1.26, 2.00, 1.25),
            (28, "Ni", "Nickel", 2.0, 3, 1.24, 1.63, 1.24),
            (29, "Cu", "Copper", 2.0, 2, 1.32, 1.40, 1.28),
            (30, "Zn", "Zinc", 2.0, 2, 1.22, 1.39, 1.34),
            (31, "Ga", "Gallium", 3.0, 3, 1.22, 1.87, 1.35),
            (32, "Ge", "Germanium", 4.0, 4, 1.20, 2.11, None),
            (33, "As", "Arsenic", 3.0, 5, 1.19, 1.85, None),
            (34, "Se", "Selenium", 2.0, 6, 1.20, 1.90, None),
            (35, "Br", "Bromine", 1.0, 7, 1.20, 1.85, None),
            (36, "Kr", "Krypton", 0.5, 8, 1.16, 2.02, None),
            (37, "Rb", "Rubidium", 1.0, 1, 2.20, 3.03, 2.48),
            (38, "Sr", "Strontium", 2.0, 2, 1.95, 2.49, 2.15),
            (39, "Y", "Yttrium", 3.0, 3, 1.90, 2.00, 1.80),
            (40, "Zr", "Zirconium", 4.0, 4, 1.75, 2.00, 1.60),
            (41, "Nb", "Niobium", 5.0, 5, 1.64, 2.00, 1.46),
            (42, "Mo", "Molybdenum", 4.0, 6, 1.54, 2.00, 1.39),
            (43, "Tc", "Technetium", 4.0, 5, 1.47, 2.00, 1.36),
            (44, "Ru", "Ruthenium", 4.0, 3, 1.46, 2.00, 1.34),
            (45, "Rh", "Rhodium", 4.0, 3, 1.42, 1.63, 1.34),
            (46, "Pd", "Palladium", 4.0, 3, 1.39, 1.72, 1.37),
            (47, "Ag", "Silver", 1.0, 2, 1.45, 1.58, 1.44),
            (48, "Cd", "Cadmium", 2.0, 2, 1.44, 1.93, 1.51),
            (49, "In", "Indium", 3.0, 3, 1.42, 2.17, 1.67),
            (50, "Sn", "Tin", 4.0, 4, 1.39, 2.06, None),
            (51, "Sb", "Antimony", 3.0, 5, 1.39, 2.06, None),
            (52, "Te", "Tellurium", 2.0, 6, 1.38, 2.06, None),
            (53, "I", "Iodine", 1.0, 7, 1.39, 1.98, None),
            (54, "Xe", "Xenon", 0.5, 8, 1.40, 2.16, None),
            (55, "Cs", "Caesium", 1.0, 1, 2.44, 3.43, 2.65),
            (56, "Ba", "Barium", 2.0, 2, 2.15, 2.68, 2.22),
            (57, "La", "Lanthanum", 3.0, 3, 2.07, 2.10, 1.87),
            (58, "Ce", "Cerium", 4.0, 3, 2.04, 2.10, 1.818),
            (59, "Pr", "Praseodymium", 3.0, 3, 2.03, 2.10, 1.824),
            (60, "Nd", "Neodymium", 3.0, 3, 2.01, 2.10, 1.814),
            (61, "Pm", "Promethium", 3.0, 3, 1.99, 2.10, 1.834),
            (62, "Sm", "Samarium", 3.0, 3, 1.98, 2.10, 1.804),
            (63, "Eu", "Europium", 3.0, 3, 1.98, 2.10, 1.804),
            (64, "Gd", "Gadolinium", 3.0, 3, 1.96, 2.10, 1.804),
            (65, "Tb", "Terbium", 3.0, 3, 1.94, 2.10, 1.773),
            (66, "Dy", "Dysprosium", 3.0, 3, 1.92, 2.10, 1.781),
            (67, "Ho", "Holmium", 3.0, 3, 1.92, 2.10, 1.762),
            (68, "Er", "Erbium", 3.0, 3, 1.89, 2.10, 1.761),
            (69, "Tm", "Thulium", 3.0, 3, 1.90, 2.10, 1.759),
            (70, "Yb", "Ytterbium", 3.0, 3, 1.87, 2.10, 1.76),
            (71, "Lu", "Lutetium", 3.0, 3, 1.87, 2.10, 1.738),
            (72, "Hf", "Hafnium", 4.0, 3, 1.75, 2.10, 1.59),
            (73, "Ta", "Tantalum", 5.0, 3, 1.70, 2.10, 1.46),
            (74, "W", "Tungsten", 4.0, 3, 1.62, 2.10, 1.39),
            (75, "Re", "Rhenium", 4.0, 3, 1.51, 2.10, 1.37),
            (76, "Os", "Osmium", 4.0, 3, 1.44, 2.10, 1.35),
            (77, "Ir", "Iridium", 4.0, 3, 1.41, 2.10, 1.355),
            (78, "Pt", "Platinum", 4.0, 3, 1.36, 1.75, 1.385),
            (79, "Au", "Gold", 1.0, 3, 1.36, 1.66, 1.44),
            (80, "Hg", "Mercury", 2.0, 3, 1.32, 1.55, 1.51),
            (81, "Tl", "Thallium", 3.0, 3, 1.45, 1.96, 1.70),
            (82, "Pb", "Lead", 4.0, 4, 1.46, 2.02, None),
            (83, "Bi", "Bismuth", 3.0, 5, 1.48, 2.07, None),
            (84, "Po", "Polonium", 2.0, 6, 1.40, 1.97, None),
            (85, "At", "Astatine", 1.0, 7, 1.50, 2.02, None),
            (86, "Rn", "Radon", 0.5, 8, 1.50, 2.20, None),
            (87, "Fr", "Francium", 1.0, 1, 2.60, 3.48, None),
            (88, "Ra", "Radium", 2.0, 2, 2.21, 2.83, None),
            (89, "Ac", "Actinium", 3.0, 3, 2.15, 2.20, None),
            (90, "Th", "Thorium", 4.0, 3, 2.06, 2.20, 1.79),
            (91, "Pa", "Protactinium", 4.0, 3, 2.00, 2.20, 1.63),
            (92, "U", "Uranium", 4.0, 3, 1.96, 2.20, 1.56),
            (93, "Np", "Neptunium", 4.0, 3, 1.90, 2.20, 1.55),
            (94, "Pu", "Plutonium", 4.0, 3, 1.87, 2.20, 1.59),
            (95, "Am", "Americium", 4.0, 3, 1.80, 2.20, 1.73),
            (96, "Cm", "Curium", 4.0, 3, 1.69, 2.20, 1.74),
            (97, "Bk", "Berkelium", 4.0, 3, None, None, 1.70),
            (98, "Cf", "Californium", 4.0, 3, None, None, 1.86),
            (99, "Es", "Einsteinium", 4.0, 3, None, None, 1.86),
            (100, "Fm", "Fermium", 4.0, 3, None, None, None),
            (101, "Md", "Mendelevium", 4.0, 3, None, None, None),
            (102, "No", "Nobelium", 4.0, 3, None, None, None),
            (103, "Lr", "Lawrencium", 4.0, 3, None, None, None),
            (104, "Rf", "Rutherfordium", 4.0, 3, None, None, None),
            (105, "Db", "Dubnium", 2.0, 3, None, None, None),
        ]
        """A list of atomic numbers, symbols, names, and other information, up
        to atomic number 105"""

        # scatter factor
        self.sf = [
            [0.493, 0.323, 0.140, 0.041, 10.511, 26.126, 3.142, 57.800, 0.003],
            [0.873, 0.631, 0.311, 0.178, 9.104, 3.357, 22.928, 0.982, 0.006],
            [1.128, 0.751, 0.618, 0.465, 3.955, 1.052, 85.391, 168.261, 0.038],
            [1.592, 1.128, 0.539, 0.703, 43.643, 1.862, 103.483, 0.542, 0.038],
            [2.055, 1.333, 1.098, 0.707, 23.219, 1.021, 60.350, 0.140, -0.193],
            [2.310, 1.020, 1.589, 0.865, 20.844, 10.208, 0.569, 51.651, 0.216],
            [12.213, 3.132, 2.013, 1.166, 0.006, 9.893, 28.997, 0.583, -11.529],
            [3.049, 2.287, 1.546, 0.867, 13.277, 5.701, 0.324, 32.909, 0.251],
            [3.539, 2.641, 1.517, 1.024, 10.283, 4.294, 0.262, 26.148, 0.278],
            [3.955, 3.112, 1.455, 1.125, 8.404, 3.426, 0.231, 21.718, 0.352],
            [4.763, 3.174, 1.267, 1.113, 3.285, 8.842, 0.314, 129.424, 0.676],
            [5.420, 2.174, 1.227, 2.307, 2.828, 79.261, 0.381, 7.194, 0.858],
            [6.420, 1.900, 1.594, 1.965, 3.039, 0.743, 31.547, 85.089, 1.115],
            [6.292, 3.035, 1.989, 1.541, 2.439, 32.334, 0.678, 81.694, 1.141],
            [6.435, 4.179, 1.780, 1.491, 1.907, 27.157, 0.526, 68.164, 1.115],
            [6.905, 5.203, 1.438, 1.586, 1.468, 22.215, 0.254, 56.172, 0.867],
            [11.460, 7.196, 6.256, 1.645, 0.010, 1.166, 18.519, 47.778, -9.557],
            [7.484, 6.772, 0.654, 1.644, 0.907, 14.841, 43.898, 33.393, 1.444],
            [8.219, 7.440, 1.052, 0.866, 12.795, 0.775, 213.187, 41.684, 1.423],
            [8.627, 7.387, 1.590, 1.021, 10.442, 0.660, 85.748, 178.437, 1.375],
            [9.189, 7.368, 1.641, 1.468, 9.021, 0.573, 136.108, 51.353, 1.333],
            [9.759, 7.356, 1.699, 1.902, 7.851, 0.500, 35.634, 116.105, 1.281],
            [10.297, 7.351, 2.070, 2.057, 6.866, 0.438, 26.894, 102.478, 1.220],
            [10.641, 7.354, 3.324, 1.492, 6.104, 0.392, 20.263, 98.740, 1.183],
            [11.282, 7.357, 3.019, 2.244, 5.341, 0.343, 17.867, 83.754, 1.090],
            [11.769, 7.357, 3.522, 2.305, 4.761, 0.307, 15.354, 76.881, 1.037],
            [12.284, 7.341, 4.003, 2.349, 4.279, 0.278, 13.536, 71.169, 1.012],
            [12.838, 7.292, 4.444, 2.380, 3.878, 0.257, 12.176, 66.342, 1.034],
            [13.338, 7.168, 5.616, 1.673, 3.583, 0.247, 11.397, 64.831, 1.191],
            [14.074, 7.032, 5.165, 2.410, 3.266, 0.233, 10.316, 58.710, 1.304],
            [15.235, 6.701, 4.359, 2.962, 3.067, 0.241, 10.781, 61.414, 1.719],
            [16.082, 6.375, 3.707, 3.683, 2.851, 0.252, 11.447, 54.763, 2.131],
            [16.672, 6.070, 3.431, 4.278, 2.635, 0.265, 12.948, 47.797, 2.531],
            [17.001, 5.820, 3.973, 4.354, 2.410, 0.273, 15.237, 43.816, 2.841],
            [17.179, 5.236, 5.638, 3.985, 2.172, 16.580, 0.261, 41.433, 2.956],
            [17.355, 6.729, 5.549, 3.537, 1.938, 16.562, 0.226, 39.397, 2.825],
            [17.178, 9.644, 5.140, 1.529, 1.789, 17.315, 0.275, 164.934, 3.487],
            [17.566, 9.818, 5.422, 2.669, 1.556, 14.099, 0.166, 132.376, 2.506],
            [17.776, 10.295, 5.726, 3.266, 1.403, 12.801, 0.261, 104.354, 1.912],
            [17.876, 10.948, 5.417, 3.657, 1.276, 11.916, 0.118, 87.663, 2.069],
            [17.614, 12.014, 4.042, 3.533, 1.189, 11.766, 0.205, 69.796, 3.756],
            [3.703, 17.236, 12.888, 3.743, 0.277, 1.096, 11.004, 61.658, 4.387],
            [19.130, 11.095, 4.649, 2.713, 0.864, 8.145, 21.571, 86.847, 5.404],
            [19.267, 12.918, 4.863, 1.568, 0.809, 8.435, 24.800, 94.293, 5.379],
            [19.296, 14.350, 4.734, 1.289, 0.752, 8.218, 25.875, 98.606, 5.328],
            [19.332, 15.502, 5.295, 0.606, 0.699, 7.989, 25.205, 76.899, 5.266],
            [19.281, 16.688, 4.805, 1.046, 0.645, 7.473, 24.660, 99.816, 5.179],
            [19.221, 17.644, 4.461, 1.603, 0.595, 6.909, 24.701, 87.482, 5.069],
            [19.162, 18.560, 4.295, 2.040, 0.548, 6.378, 25.850, 92.803, 4.939],
            [19.189, 19.101, 4.458, 2.466, 5.830, 0.503, 26.891, 83.957, 4.782],
            [19.642, 19.045, 5.037, 2.683, 5.303, 0.461, 27.907, 75.283, 4.591],
            [19.964, 19.014, 6.145, 2.524, 4.817, 0.421, 28.528, 70.840, 4.352],
            [20.147, 18.995, 7.514, 2.273, 4.347, 0.381, 27.766, 66.878, 4.071],
            [20.293, 19.030, 8.977, 1.990, 3.928, 0.344, 26.466, 64.266, 3.712],
            [20.389, 19.106, 10.662, 1.495, 3.569, 0.311, 24.388, 213.904, 3.335],
            [20.336, 19.297, 10.888, 2.696, 3.216, 0.276, 20.207, 167.202, 2.773],
            [20.578, 19.599, 11.373, 3.287, 2.948, 0.244, 18.773, 133.124, 2.147],
            [21.167, 19.770, 11.851, 3.330, 2.812, 0.227, 17.608, 127.113, 1.863],
            [22.044, 19.670, 12.386, 2.824, 2.774, 0.222, 16.767, 143.644, 2.058],
            [22.684, 19.685, 12.774, 2.851, 2.662, 0.211, 15.885, 137.903, 1.985],
            [23.340, 19.610, 13.123, 2.875, 2.563, 0.202, 15.101, 132.721, 2.029],
            [24.004, 19.426, 13.440, 2.896, 2.473, 0.196, 14.400, 128.007, 2.210],
            [24.627, 19.089, 13.760, 2.293, 2.388, 0.194, 13.755, 123.174, 2.575],
            [25.071, 19.080, 13.852, 3.545, 2.253, 0.182, 12.933, 101.398, 2.420],
            [25.898, 18.219, 14.317, 2.954, 2.243, 0.196, 12.665, 115.362, 3.583],
            [26.507, 17.638, 14.560, 2.966, 2.180, 0.202, 12.190, 111.874, 4.297],
            [26.905, 17.294, 14.558, 3.638, 2.071, 0.198, 11.441, 92.657, 4.568],
            [27.656, 16.428, 14.978, 2.982, 2.074, 0.224, 11.360, 105.703, 5.920],
            [28.182, 15.885, 15.154, 2.987, 2.029, 0.239, 10.998, 102.961, 6.756],
            [28.664, 15.434, 15.309, 2.990, 1.989, 0.257, 10.665, 100.417, 7.567],
            [28.948, 15.221, 15.100, 3.716, 1.902, 9.985, 0.261, 84.330, 7.976],
            [29.144, 15.173, 14.759, 4.300, 1.833, 9.600, 0.275, 72.029, 8.582],
            [29.202, 15.229, 14.514, 4.765, 1.773, 9.370, 0.296, 63.364, 9.244],
            [0.000, 0.000, 0.000, 0.000, 1.722, 9.231, 0.323, 57.725, 9.858],
            [28.762, 15.719, 14.556, 5.442, 1.672, 9.092, 0.350, 52.086, 10.472],
            [28.189, 16.155, 14.931, 5.676, 1.629, 8.979, 0.383, 48.165, 11.000],
            [27.305, 16.730, 15.611, 5.834, 1.593, 8.866, 0.418, 45.001, 11.472],
            [27.006, 17.764, 15.713, 5.784, 1.513, 8.812, 0.425, 38.610, 11.688],
            [16.882, 18.591, 25.558, 5.860, 0.461, 8.622, 1.483, 36.396, 12.066],
            [20.681, 19.042, 21.657, 5.968, 0.545, 8.448, 1.573, 38.325, 12.609],
            [27.545, 19.158, 15.538, 5.526, 0.655, 8.708, 1.963, 45.815, 13.175],
            [31.062, 13.064, 18.442, 5.970, 0.690, 2.358, 8.618, 47.258, 13.412],
            [33.369, 12.951, 16.588, 6.469, 0.704, 2.924, 8.794, 48.009, 13.578],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
            [35.660, 23.103, 12.598, 4.087, 0.589, 3.652, 18.599, 117.020, 13.527],
            [35.564, 23.422, 12.747, 4.807, 0.563, 3.462, 17.831, 99.172, 13.431],
            [35.885, 23.295, 14.189, 4.173, 0.548, 3.415, 16.924, 105.251, 13.429],
            [0.000, 0.000, 0.000, 0.000, 0.530, 3.335, 16.143, 101.371, 13.393],
            [36.187, 23.596, 15.640, 4.186, 0.512, 3.254, 15.362, 97.491, 13.357],
            [36.526, 23.808, 16.771, 3.479, 0.499, 3.264, 14.946, 105.980, 13.381],
        ]
        """A list of scatter factors for the elements"""

        self.z = None
        """atomic number"""
        self.short_name = None
        """atomic symbol"""
        self.long_name = None
        """atomic name"""
        self.valence = None
        """valence value"""
        self.valence_electrons = None
        """number of valence electrons"""
        self.covalent_radius = None
        """atomic radius used for distance checking within crystals"""
        self.vdw_radius = None
        """atomic radius used for volume estimation within crystals"""
        self.metallic_radius = None
        """atomic radius used for distance checking within metallic crystals"""

        pos = None

        try:
            int(self.input)
            self.z = self.input

            for i, el in enumerate(self.elements_list):
                if el[0] == self.z:
                    pos = i
                    self.short_name = el[1]
                    self.long_name = el[2]
                    break
        except ValueError:
            self.short_name = self.input
            for i, el in enumerate(self.elements_list):
                if el[1] == self.short_name:
                    pos = i
                    self.z = el[0]
                    self.long_name = el[2]
                    break

            if not self.z:
                self.short_name = None
                self.long_name = self.input
                for i, el in enumerate(self.elements_list):
                    if el[2] == self.long_name:
                        pos = i
                        self.z = el[0]
                        self.short_name = el[1]
                        break
                if not self.z:
                    self.long_name = None

        if pos is not None:
            self.valence = self.elements_list[pos][3]
            self.valence_electrons = self.elements_list[pos][4]
            self.covalent_radius = self.elements_list[pos][5]
            self.vdw_radius = self.elements_list[pos][6]
            self.metallic_radius = self.elements_list[pos][7]
            self.scatter = self.sf[pos]

    def get_all(self, pos):
        """
        Return all [pos] elements in the full element list

        Args:
            pos: the index of the elements to retrieve

        Returns:
            a list containing only the [pos] elements of self.elements_list
        """
        els = []
        for el in self.elements_list:
            els.append(el[pos])
        return els

    def get_sf(self, pos):
        """
        Get the scattering factor for an element.

        Args:
            pos: the atomic number of the element

        Returns:
            the scattering factor of the element
        """
        els = []
        for el in self.sf:
            els.append(el[pos])
        return els

    # TODO: add docstrings for miscellaneous functions
    def all_z(self):
        return self.get_all(0)

    def all_short_names(self):
        return self.get_all(1)

    def all_long_names(self):
        return self.get_all(2)

    def all_valences(self):
        return self.get_all(3)

    def all_valence_electrons(self):
        return self.get_all(4)

    def all_covalent_radii(self):
        return self.get_all(5)

    def all_vdw_radii(self):
        return self.get_all(6)

    def all_metallic_radii(self):
        return self.get_all(7)

    def get_sf(self):
        return self.get_sf()

    def number_from_specie(specie):
        if type(specie) == int or type(specie) == float:
            if specie <= 105 and specie >= 1:
                index = int(specie)
            else:
                print(specie)
                print("Error: Atomic number must be between 1 and 105.")
                return
        elif type(specie) == str:
            try:
                el = Element(specie)
                index = el.z
            except:
                print("Error: Invalid atomic symbol, name, or number.")
                return
        elif type(specie) == Element:
            try:
                index = specie.z
            except:
                print("Error: Element object has no atomic number 'z'.")
                return
        else:
            try:
                el = Element(specie.number)
                index = el.z
            except:
                print("Error: Invalid atomic symbol, name, or number.")
                return
        return index


"""Utilities for coordinate-based meta-analysis estimators."""
import logging
import os

import nibabel as nib
import numpy as np
import numpy.linalg as npl
from scipy import ndimage

from .. import references
from ..due import due
from ..extract import download_peaks2maps_model
from ..utils import _determine_chunk_size

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
LGR = logging.getLogger(__name__)

def model_fn(features, labels, mode, params):
    """Run model function used internally by peaks2maps.

    .. deprecated:: 0.0.11
        `model_fn` will be removed in NiMARE 0.0.13.

    .. versionadded:: 0.0.4

    """
    import tensorflow as tf
    from tensorflow.python.estimator.export.export_output import PredictOutput

    ngf = 64
    layers = []

    training_flag = mode == tf.estimator.ModeKeys.TRAIN

    input_images_placeholder = tf.expand_dims(features, -1)

    conv_args = {
        "strides": 2,
        "kernel_size": 4,
        "padding": "valid",
        "activation": tf.nn.leaky_relu,
        "kernel_initializer": tf.random_normal_initializer(0, 0.02),
        "name": "conv",
        "use_bias": False,
    }

    deconv_args = conv_args.copy()
    deconv_args["padding"] = "same"
    deconv_args["name"] = "deconv"

    batchnorm_args = {
        "scale": True,
        "gamma_initializer": tf.random_normal_initializer(1.0, 0.02),
        "center": True,
        "beta_initializer": tf.zeros_initializer(),
        "name": "batchnorm",
        "training": training_flag,
    }

    def pad_and_conv(input, out_channels, conv_args):
        padded_input = tf.pad(input, [[0, 0], [1, 1], [1, 1], [1, 1], [0, 0]], mode="CONSTANT")
        convolved = tf.compat.v1.layers.conv3d(padded_input, out_channels, **conv_args)
        return convolved

    # encoder_1: [batch, 256, 256, in_channels] => [batch, 128, 128, ngf]
    with tf.compat.v1.variable_scope("encoder_1"):
        this_args = conv_args.copy()
        output = pad_and_conv(input_images_placeholder, ngf, this_args)
        layers.append(output)

    layer_specs = [
        (ngf * 2, 0.2),
        # encoder_3: [batch, 64, 64, ngf * 2] => [batch, 32, 32, ngf * 4]
        (ngf * 2, 0.2),
        # encoder_3: [batch, 64, 64, ngf * 2] => [batch, 32, 32, ngf * 4]
        (ngf * 4, 0.2),
        # encoder_4: [batch, 32, 32, ngf * 4] => [batch, 16, 16, ngf * 8]
        (ngf * 8, 0.2),
        # encoder_5: [batch, 16, 16, ngf * 8] => [batch, 8, 8, ngf * 8]
        # ngf * 8,
        # # encoder_6: [batch, 8, 8, ngf * 8] => [batch, 4, 4, ngf * 8]
    ]

    for out_channels, dropout in layer_specs:
        with tf.compat.v1.variable_scope("encoder_%d" % (len(layers) + 1)):
            # [batch, in_height, in_width, in_channels] => [batch, in_height/2, in_width/2,
            # out_channels]
            convolved = pad_and_conv(layers[-1], out_channels, conv_args)
            output = tf.compat.v1.layers.batch_normalization(convolved, **batchnorm_args)
            if dropout > 0.0:
                output = tf.compat.v1.layers.dropout(output, rate=dropout, training=training_flag)
            layers.append(output)

    layer_specs = [
        # (ngf * 8, 0.5),
        # # decoder_6: [batch, 4, 4, ngf * 8 * 2] => [batch, 8, 8, ngf * 8 * 2]
        (ngf * 8, 0.2),
        # decoder_5: [batch, 8, 8, ngf * 8 * 2] => [batch, 16, 16, ngf * 8 * 2]
        (ngf * 4, 0.2),
        # decoder_4: [batch, 16, 16, ngf * 8 * 2] => [batch, 32, 32, ngf * 4 * 2]
        (ngf * 2, 0.2),
        # decoder_3: [batch, 32, 32, ngf * 4 * 2] => [batch, 64, 64, ngf * 2 * 2]
        (ngf * 2, 0.2),
        # decoder_2: [batch, 64, 64, ngf * 2 * 2] => [batch, 128, 128, ngf * 2]
    ]

    num_encoder_layers = len(layers)
    for decoder_layer, (out_channels, dropout) in enumerate(layer_specs):
        skip_layer = num_encoder_layers - decoder_layer - 1
        with tf.compat.v1.variable_scope("decoder_%d" % (skip_layer + 1)):
            if decoder_layer == 0:
                # first decoder layer doesn't have skip connections
                # since it is directly connected to the skip_layer
                input = layers[-1]
            else:
                input = tf.concat([layers[-1], layers[skip_layer]], axis=4)

            output = tf.compat.v1.layers.conv3d_transpose(input, out_channels, **deconv_args)
            output = tf.compat.v1.layers.batch_normalization(output, **batchnorm_args)

            if dropout > 0.0:
                output = tf.compat.v1.layers.dropout(output, rate=dropout, training=training_flag)

            layers.append(output)

    # decoder_1: [batch, 128, 128, ngf * 2] => [batch, 256, 256, generator_outputs_channels]
    with tf.compat.v1.variable_scope("decoder_1"):
        input = tf.concat([layers[-1], layers[0]], axis=4)
        this_args = deconv_args.copy()
        this_args["activation"] = None
        output = tf.compat.v1.layers.conv3d_transpose(input, 1, **this_args)
        layers.append(output)

    predictions = tf.squeeze(layers[-1], -1)

    if mode == tf.estimator.ModeKeys.PREDICT:
        temp = tf.compat.v1.saved_model.signature_constants
        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=predictions,
            export_outputs={temp.DEFAULT_SERVING_SIGNATURE_DEF_KEY: PredictOutput(predictions)},
        )
    else:
        labels, filenames = labels
        loss = tf.losses.mean_squared_error(labels, predictions)

        # Add a scalar summary for the snapshot loss.
        # Create the gradient descent optimizer with the given learning rate.
        optimizer = tf.train.AdamOptimizer(params.learning_rate)
        # Use the optimizer to apply the gradients that minimize the loss
        # (and also increment the global step counter) as a single training step.
        extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(extra_update_ops):
            train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())

        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=predictions,
            loss=loss,
            train_op=train_op,
            export_outputs={"output": predictions},
        )

def _get_resize_arg(target_shape):
    """Get resizing arguments, as used by peaks2maps.

    .. deprecated:: 0.0.11
        `_get_resize_arg` will be removed in NiMARE 0.0.13.

    .. versionadded:: 0.0.1

    """
    mni_shape_mm = np.array([148.0, 184.0, 156.0])
    target_resolution_mm = np.ceil(mni_shape_mm / np.array(target_shape)).astype(np.int32)
    target_affine = np.array(
        [
            [4.0, 0.0, 0.0, -75.0],
            [0.0, 4.0, 0.0, -105.0],
            [0.0, 0.0, 4.0, -70.0],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    target_affine[0, 0] = target_resolution_mm[0]
    target_affine[1, 1] = target_resolution_mm[1]
    target_affine[2, 2] = target_resolution_mm[2]
    return target_affine, list(target_shape)

def _get_generator(contrasts_coordinates, target_shape, affine, skip_out_of_bounds=False):
    """Get generator, as used by peaks2maps."""

    def generator():
        for contrast in contrasts_coordinates:
            encoded_coords = np.zeros(list(target_shape))
            for real_pt in contrast:
                vox_pt = np.rint(nib.affines.apply_affine(npl.inv(affine), real_pt)).astype(int)
                if skip_out_of_bounds and (vox_pt[0] >= 32 or vox_pt[1] >= 32 or vox_pt[2] >= 32):
                    continue
                encoded_coords[vox_pt[0], vox_pt[1], vox_pt[2]] = 1
            yield (encoded_coords, encoded_coords)

    return generator

@due.dcite(
    references.PEAKS2MAPS,
    description="Transforms coordinates of peaks to unthresholded maps using a deep "
    "convolutional neural net.",
)
def compute_p2m_ma(
    contrasts_coordinates, skip_out_of_bounds=True, tf_verbosity_level=None, model_dir="auto"
):
    """Generate modeled activation (MA) maps using deep ConvNet model peaks2maps.

    .. deprecated:: 0.0.11
        `compute_p2m_ma` will be removed in NiMARE 0.0.13.

    Parameters
    ----------
    contrasts_coordinates : list of lists that are len == 3
        List of contrasts and their coordinates
    skip_out_of_bounds : bool, optional
        Remove coordinates outside of the bounding box of the peaks2maps model
    tf_verbosity_level : int
        Tensorflow verbosity logging level
    model_dir : str, optional
        Location of peaks2maps model. Default is "auto".

    Returns
    -------
    ma_values : array-like
        1d array of modeled activation values.
    """
    try:
        import tensorflow as tf
    except ImportError as e:
        if "No module named 'tensorflow'" in str(e):
            raise Exception(
                "tensorflow not installed - see https://www.tensorflow.org/install/ "
                "for instructions"
            )
        else:
            raise

    if tf_verbosity_level is None:
        tf_verbosity_level = tf.compat.v1.logging.FATAL
    target_shape = (32, 32, 32)
    affine, _ = _get_resize_arg(target_shape)
    tf.compat.v1.logging.set_verbosity(tf_verbosity_level)

    def generate_input_fn():
        dataset = tf.compat.v1.data.Dataset.from_generator(
            _get_generator(
                contrasts_coordinates, target_shape, affine, skip_out_of_bounds=skip_out_of_bounds
            ),
            (tf.float32, tf.float32),
            (tf.TensorShape(target_shape), tf.TensorShape(target_shape)),
        )
        dataset = dataset.batch(1)
        iterator = dataset.make_one_shot_iterator()
        return iterator.get_next()

    # download_peaks2maps_model expects None for "auto"
    model_dir = None if model_dir == "auto" else model_dir
    model_dir = download_peaks2maps_model(data_dir=model_dir)
    model = tf.estimator.Estimator(model_fn, model_dir=model_dir)

    results = model.predict(generate_input_fn)
    results = [result for result in results]
    assert len(results) == len(contrasts_coordinates), "returned %d" % len(results)

    niis = [nib.Nifti1Image(np.squeeze(result), affine) for result in results]
    return niis

def compute_kda_ma(
    shape,
    vox_dims,
    ijks,
    r,
    value=1.0,
    exp_idx=None,
    sum_overlap=False,
    memory_limit=None,
    memmap_filename=None,
):
    """Compute (M)KDA modeled activation (MA) map.

    .. versionchanged:: 0.0.8

        * [ENH] Add *memmap_filename* parameter for memory mapping arrays.

    .. versionadded:: 0.0.4

    Replaces the values around each focus in ijk with binary sphere.

    Parameters
    ----------
    shape : :obj:`tuple`
        Shape of brain image + buffer. Typically (91, 109, 91).
    vox_dims : array_like
        Size (in mm) of each dimension of a voxel.
    ijks : array-like
        Indices of foci. Each row is a coordinate, with the three columns
        corresponding to index in each of three dimensions.
    r : :obj:`int`
        Sphere radius, in mm.
    value : :obj:`int`
        Value for sphere.
    exp_idx : array_like
        Optional indices of experiments. If passed, must be of same length as
        ijks. Each unique value identifies all coordinates in ijk that come from
        the same experiment. If None passed, it is assumed that all coordinates
        come from the same experiment.
    sum_overlap : :obj:`bool`
        Whether to sum voxel values in overlapping spheres.
    memory_limit : :obj:`str` or None, optional
        Memory limit to apply to data. If None, no memory management will be applied.
        Otherwise, the memory limit will be used to (1) assign memory-mapped files and
        (2) restrict memory during array creation to the limit.
        Default is None.
    memmap_filename : :obj:`str`, optional
        If passed, use this file for memory mapping arrays

    Returns
    -------
    kernel_data : :obj:`numpy.array`
        3d or 4d array. If `exp_idx` is none, a 3d array in the same shape as
        the `shape` argument is returned. If `exp_idx` is passed, a 4d array
        is returned, where the first dimension has size equal to the number of
        unique experiments, and the remaining 3 dimensions are equal to `shape`.
    """
    squeeze = exp_idx is None
    if exp_idx is None:
        exp_idx = np.ones(len(ijks))

    uniq, exp_idx = np.unique(exp_idx, return_inverse=True)
    n_studies = len(uniq)

    kernel_shape = (n_studies,) + shape
    if memmap_filename:
        # Use a memmapped 4D array
        kernel_data = np.memmap(memmap_filename, dtype=type(value), mode="w+", shape=kernel_shape)
    else:
        kernel_data = np.zeros(kernel_shape, dtype=type(value))

    n_dim = ijks.shape[1]
    xx, yy, zz = [slice(-r // vox_dims[i], r // vox_dims[i] + 0.01, 1) for i in range(n_dim)]
    cube = np.vstack([row.ravel() for row in np.mgrid[xx, yy, zz]])
    kernel = cube[:, np.sum(np.dot(np.diag(vox_dims), cube) ** 2, 0) ** 0.5 <= r]

    if memory_limit:
        chunk_size = _determine_chunk_size(limit=memory_limit, arr=ijks[0])

    for i, peak in enumerate(ijks):
        sphere = np.round(kernel.T + peak)
        idx = (np.min(sphere, 1) >= 0) & (np.max(np.subtract(sphere, shape), 1) <= -1)
        sphere = sphere[idx, :].astype(int)
        exp = exp_idx[i]
        if sum_overlap:
            kernel_data[exp][tuple(sphere.T)] += value
        else:
            kernel_data[exp][tuple(sphere.T)] = value

        if memmap_filename and i % chunk_size == 0:
            # Write changes to disk
            kernel_data.flush()

    if squeeze:
        kernel_data = np.squeeze(kernel_data, axis=0)

    return kernel_data

def compute_ale_ma(shape, ijk, kernel):
    """Generate ALE modeled activation (MA) maps.

    Replaces the values around each focus in ijk with the contrast-specific
    kernel. Takes the element-wise maximum when looping through foci, which
    accounts for foci which are near to one another and may have overlapping
    kernels.

    Parameters
    ----------
    shape : tuple
        Shape of brain image + buffer. Typically (91, 109, 91) + (30, 30, 30).
    ijk : array-like
        Indices of foci. Each row is a coordinate, with the three columns
        corresponding to index in each of three dimensions.
    kernel : array-like
        3D array of smoothing kernel. Typically of shape (30, 30, 30).

    Returns
    -------
    ma_values : array-like
        1d array of modeled activation values.
    """
    ma_values = np.zeros(shape)
    mid = int(np.floor(kernel.shape[0] / 2.0))
    mid1 = mid + 1
    for j_peak in range(ijk.shape[0]):
        i, j, k = ijk[j_peak, :]
        xl = max(i - mid, 0)
        xh = min(i + mid1, ma_values.shape[0])
        yl = max(j - mid, 0)
        yh = min(j + mid1, ma_values.shape[1])
        zl = max(k - mid, 0)
        zh = min(k + mid1, ma_values.shape[2])
        xlk = mid - (i - xl)
        xhk = mid - (i - xh)
        ylk = mid - (j - yl)
        yhk = mid - (j - yh)
        zlk = mid - (k - zl)
        zhk = mid - (k - zh)

        if (
            (xl >= 0)
            & (xh >= 0)
            & (yl >= 0)
            & (yh >= 0)
            & (zl >= 0)
            & (zh >= 0)
            & (xlk >= 0)
            & (xhk >= 0)
            & (ylk >= 0)
            & (yhk >= 0)
            & (zlk >= 0)
            & (zhk >= 0)
        ):
            ma_values[xl:xh, yl:yh, zl:zh] = np.maximum(
                ma_values[xl:xh, yl:yh, zl:zh], kernel[xlk:xhk, ylk:yhk, zlk:zhk]
            )
    return ma_values

@due.dcite(references.ALE_KERNEL, description="Introduces sample size-dependent kernels to ALE.")
def get_ale_kernel(img, sample_size=None, fwhm=None):
    """Estimate 3D Gaussian and sigma (in voxels) for ALE kernel given sample size or fwhm."""
    if sample_size is not None and fwhm is not None:
        raise ValueError('Only one of "sample_size" and "fwhm" may be specified')
    elif sample_size is None and fwhm is None:
        raise ValueError('Either "sample_size" or "fwhm" must be provided')
    elif sample_size is not None:
        uncertain_templates = (
            5.7 / (2.0 * np.sqrt(2.0 / np.pi)) * np.sqrt(8.0 * np.log(2.0))
        )  # pylint: disable=no-member
        # Assuming 11.6 mm ED between matching points
        uncertain_subjects = (11.6 / (2 * np.sqrt(2 / np.pi)) * np.sqrt(8 * np.log(2))) / np.sqrt(
            sample_size
        )  # pylint: disable=no-member
        fwhm = np.sqrt(uncertain_subjects ** 2 + uncertain_templates ** 2)

    fwhm_vox = fwhm / np.sqrt(np.prod(img.header.get_zooms()))
    sigma_vox = (
        fwhm_vox * np.sqrt(2.0) / (np.sqrt(2.0 * np.log(2.0)) * 2.0)
    )  # pylint: disable=no-member

    data = np.zeros((31, 31, 31))
    mid = int(np.floor(data.shape[0] / 2.0))
    data[mid, mid, mid] = 1.0
    kernel = ndimage.filters.gaussian_filter(data, sigma_vox, mode="constant")

    # Crop kernel to drop surrounding zeros
    mn = np.min(np.where(kernel > np.spacing(1))[0])
    mx = np.max(np.where(kernel > np.spacing(1))[0])
    kernel = kernel[mn : mx + 1, mn : mx + 1, mn : mx + 1]
    mid = int(np.floor(data.shape[0] / 2.0))
    return sigma_vox, kernel


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

"""hyperscreen is an improved background rejection algorithm
for """

from __future__ import division
from __future__ import print_function

import colorama
from tqdm import tqdm as progressbar
import sys
import os

import warnings

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

from astropy.io import fits
from astropy.table import Table

import skimage.data as data
import skimage.segmentation as seg
import skimage.filters as filters
import skimage.draw as draw
import skimage.color as color

import pandas as pd
import numpy as np
np.seterr(divide='ignore')

colorama.init()

class HRCevt1:
    """This is a conceptual class representation of a Chandra High Resolution Camera (HRC) Level 1 Event File

    :return: HRCevt1 object
    :rtype: pandas.DataFrame or astropy.table.table.Table
    """

    def __init__(self, evt1file, verbose=False, as_astropy_table=False):
        """The constructor method for the HRCevt1 class

        :param evt1file: A .fits (or fits.gz) file containing the level 1 event list. If downloaded from the Chandra database, this file always has a *evt1.fits extension. This event list includes all events telemetered.
        :type evt1file: .fits or .fits.gz
        :param verbose: Set verbose=True to make the constructor chatty on the command line, defaults to False
        :type verbose: bool, optional
        :param as_astropy_table: Set as_astropy_table to True in order to have the HRCevt1 constructor method return an Astropy Table object, rather than a Pandas DataFrame. Defaults to False.
        :type as_astropy_table: bool, optional
        """

        # Define how chatty to be
        self.verbose = verbose

        if self.verbose is True:
            print(colorama.Fore.BLUE + '\nParsing HRC EVT1 file...', end=" ")
        # Do a standard read in of the EVT1 fits table
        self.filename = evt1file
        self.hdulist = fits.open(evt1file)
        self.data = Table(self.hdulist[1].data)
        self.header = self.hdulist[1].header
        self.gti = self.hdulist[2].data
        self.hdulist.close()  # Don't forget to close your fits file!

        # Make sure the user isn't running this on an ACIS observation!
        if self.header["DETNAM"][:4] == 'ACIS':
            raise Exception(
                "ERROR: HRCevt1 objects can only be initialized for Chandra/HRC observations. This is a Chandra/ACIS observation.")

        # Populate the fp, fb values for ever event
        if self.verbose is True:
            print(colorama.Fore.BLUE + 'Calculating fp, fb values...', end=" ")
        fp_u, fb_u, fp_v, fb_v = self.calculate_fp_fb()

        # Populate the fp, fb values for ever event
        if self.verbose is True:
            print(colorama.Fore.BLUE + 'Applying GTI mask... ', end=" ")
        self.gti.starts = self.gti['START']
        self.gti.stops = self.gti['STOP']

        self.gtimask = (self.data["time"] > self.gti.starts[0]) & (
            self.data["time"] < self.gti.stops[-1])

        # Populate the fp, fb values for every event
        if self.verbose is True:
            print(colorama.Fore.BLUE + 'Populating metadata columns...', end=" ")
        self.data["fp_u"] = fp_u
        self.data["fb_u"] = fb_u
        self.data["fp_v"] = fp_v
        self.data["fb_v"] = fb_v

        # Make individual status bit columns with legible names
        self.data["AV3 corrected for ringing"] = self.data["status"][:, 0]
        self.data["AU3 corrected for ringing"] = self.data["status"][:, 1]
        self.data["Event impacted by prior event (piled up)"] = self.data["status"][:, 2]
        # Bit 4 (Python 3) is spare
        self.data["Shifted event time"] = self.data["status"][:, 4]
        self.data["Event telemetered in NIL mode"] = self.data["status"][:, 5]
        self.data["V axis not triggered"] = self.data["status"][:, 6]
        self.data["U axis not triggered"] = self.data["status"][:, 7]
        self.data["V axis center blank event"] = self.data["status"][:, 8]
        self.data["U axis center blank event"] = self.data["status"][:, 9]
        self.data["V axis width exceeded"] = self.data["status"][:, 10]
        self.data["U axis width exceeded"] = self.data["status"][:, 11]
        self.data["Shield PMT active"] = self.data["status"][:, 12]
        # Bit 14 (Python 13) is hardware spare
        self.data["Upper level discriminator not exceeded"] = self.data["status"][:, 14]
        self.data["Lower level discriminator not exceeded"] = self.data["status"][:, 15]
        self.data["Event in bad region"] = self.data["status"][:, 16]
        self.data["Amp total on V or U = 0"] = self.data["status"][:, 17]
        self.data["Incorrect V center"] = self.data["status"][:, 18]
        self.data["Incorrect U center"] = self.data["status"][:, 19]
        self.data["PHA ratio test failed"] = self.data["status"][:, 20]
        self.data["Sum of 6 taps = 0"] = self.data["status"][:, 21]
        self.data["Grid ratio test failed"] = self.data["status"][:, 22]
        self.data["ADC sum on V or U = 0"] = self.data["status"][:, 23]
        self.data["PI exceeding 255"] = self.data["status"][:, 24]
        self.data["Event time tag is out of sequence"] = self.data["status"][:, 25]
        self.data["V amp flatness test failed"] = self.data["status"][:, 26]
        self.data["U amp flatness test failed"] = self.data["status"][:, 27]
        self.data["V amp saturation test failed"] = self.data["status"][:, 28]
        self.data["U amp saturation test failed"] = self.data["status"][:, 29]
        self.data["V hyperbolic test failed"] = self.data["status"][:, 30]
        self.data["U hyperbolic test failed"] = self.data["status"][:, 31]
        self.data["Hyperbola test passed"] = np.logical_not(np.logical_or(
            self.data['U hyperbolic test failed'], self.data['V hyperbolic test failed']))
        self.data["Hyperbola test failed"] = np.logical_or(
            self.data['U hyperbolic test failed'], self.data['V hyperbolic test failed'])

        self.obsid = self.header["OBS_ID"]
        self.obs_date = self.header["DATE"]
        self.target = self.header["OBJECT"]
        self.detector = self.header["DETNAM"]
        self.grating = self.header["GRATING"]
        self.exptime = self.header["EXPOSURE"]

        self.numevents = len(self.data["time"])
        self.goodtimeevents = len(self.data["time"][self.gtimask])
        self.badtimeevents = self.numevents - self.goodtimeevents

        self.hyperbola_passes = np.sum(np.logical_or(
            self.data['U hyperbolic test failed'], self.data['V hyperbolic test failed']))
        self.hyperbola_failures = np.sum(np.logical_not(np.logical_or(
            self.data['U hyperbolic test failed'], self.data['V hyperbolic test failed'])))

        if self.hyperbola_passes + self.hyperbola_failures != self.numevents:
            warnings.warn("Number of Hyperbola Test Failures and Passes ({}) does not equal total number of events ({}).".format(
                self.hyperbola_passes + self.hyperbola_failures, self.numevents))

        if self.verbose is True:
            print(colorama.Fore.GREEN + 'Done')

        if self.verbose is True:
            if as_astropy_table is True:
                read_type = "Astropy Table"
            else:
                read_type = "Pandas DataFrame"
            print(colorama.Fore.CYAN + 'Observation Details: ')
            print(colorama.Fore.CYAN + 'ObsID {}  |    {}    |    {}      |    {} ksec     |      {:,} counts (level 1 events)'.format(self.obsid,
                                                                                self.target, self.detector, np.round(self.exptime/1000, 2), self.numevents))
            print(colorama.Fore.RED + '\nConverting EVT1 file to {}...'.format(read_type), end=" ")

        if as_astropy_table is False:
            # Multidimensional columns don't grok with Pandas
            self.data.remove_column('status')
            self.data = self.data.to_pandas()

        if self.verbose is True:
            print(colorama.Fore.GREEN + 'Done')

    def __str__(self):
        """This method returns the string representation of the HRCevt1 object. It is called when the print() or str() function is invoked on an HRCevt1 object.

        :return: A string describing the HRCevt1 object
        :rtype: str
        """
        return "HRC EVT1 object with {} events. Data is packaged as a Pandas Dataframe (or an Astropy Table if as_astropy_table=True on initialization.)".format(self.numevents)

    def calculate_fp_fb(self):
        """Method to calculate the Fine Position (f_p) and normalized central tap amplitude (fb) for the HRC U- and V- axes.

        :return: fp_u, fb_u, fp_v, fb_v; the calculated fine positions and normalized central tap amplitudes, respectively, for the HRC U- and V- axes of the I or S detector
        :rtype: float
        """

        a_u = self.data["au1"]  # otherwise known as "a1"
        b_u = self.data["au2"]  # "a2"
        c_u = self.data["au3"]  # "a3"

        a_v = self.data["av1"]
        b_v = self.data["av2"]
        c_v = self.data["av3"]

        with np.errstate(invalid='ignore'):
            # Do the U axis
            fp_u = ((c_u - a_u) / (a_u + b_u + c_u))
            fb_u = b_u / (a_u + b_u + c_u)

            # Do the V axis
            fp_v = ((c_v - a_v) / (a_v + b_v + c_v))
            fb_v = b_v / (a_v + b_v + c_v)

        return fp_u, fb_u, fp_v, fb_v

    def threshold(self, img, bins, softening=None):
        """HyperScreen (a) separates events by both axis and tap, (b) creates an
        image (a 2D histogram, with bin sizes dependent on the total number of counts
        in the observation), then (c) performs efficient image segmentation on that image
        by applying Otsu's Method (https://en.wikipedia.org/wiki/Otsu%27s_method) to
        create a bespoke threshold for each tap image. This efficiently separates the 'boomerang'
        locus for 'real' events with the 'everything else' region for background events. This function
        performs operation (c).

        Arguments:
            img {[type]} -- [description]
            bins {[type]} -- [description]
        """

        # You don't want to be verbose in this function; it's called many times

        thresh_img = img.copy()
        thresh_img[img == 0] = np.nan

        if softening is None:
            thresh = filters.threshold_otsu(img)
        elif isinstance(softening, float):
            otsu_thresh = filters.threshold_otsu(img)
            thresh = otsu_thresh - (otsu_thresh * softening)

        # "If you don't ignore the warning, you'll get a warning" ~~ G. Tremblay
        with np.errstate(invalid='ignore'):
            thresh_img[thresh_img < thresh] = np.nan
        thresh_img[:int(bins[1] / 2), :] = np.nan
    #     thresh_img[:,int(bins[1]-5):] = np.nan

        return thresh_img

    def hyperscreen(self, softening=1.0):
        """[summary]

        Returns:
            [type] -- [description]
        """

        data = self.data[self.data['Hyperbola test passed']]

        # taprange = range(data['crsu'].min(), data['crsu'].max() + 1)
        taprange_u = range(data['crsu'].min() - 1, data['crsu'].max() + 1)
        taprange_v = range(data['crsv'].min() - 1, data['crsv'].max() + 1)

        if self.numevents < 100000:
            bins = [50, 50]  # number of bins
        else:
            bins = [200, 200]

        # Instantiate these empty dictionaries to hold our results
        u_axis_survivals = {}
        v_axis_survivals = {}

        if self.verbose is False:
            progressbar_disable = True
        elif self.verbose is True:
            progressbar_disable = False

        if self.verbose is True:
            print(colorama.Fore.YELLOW + "\nApplying Otsu's Method to every Tap-specific boomerang across U-axis taps {} through {}".format(taprange_u[0] + 1, taprange_u[-1] + 1))

        skiptaps_u = []
        skiptaps_v = []

        for tap in progressbar(taprange_u, disable=progressbar_disable, ascii=False):
            # Do the U axis
            tapmask_u = data[data['crsu'] == tap].index.values
            if len(tapmask_u) < 20:
                skiptaps_u.append((tap + 1, len(tapmask_u)))
                continue
            keep_u = np.isfinite(data['fb_u'][tapmask_u])

            hist_u, xbounds_u, ybounds_u = np.histogram2d(
                data['fb_u'][tapmask_u][keep_u], data['fp_u'][tapmask_u][keep_u], bins=bins)
            thresh_hist_u = self.threshold(
                hist_u, bins=bins, softening=softening)

            posx_u = np.digitize(data['fb_u'][tapmask_u], xbounds_u)
            posy_u = np.digitize(data['fp_u'][tapmask_u], ybounds_u)
            hist_mask_u = (posx_u > 0) & (posx_u <= bins[0]) & (
                posy_u > -1) & (posy_u <= bins[1])

            # Values of the histogram where the points are
            hhsub_u = thresh_hist_u[posx_u[hist_mask_u] -
                                    1, posy_u[hist_mask_u] - 1]
            pass_fb_u = data['fb_u'][tapmask_u][hist_mask_u][np.isfinite(
                hhsub_u)]

            u_axis_survivals["U Axis Tap {:02d}".format(
                tap)] = pass_fb_u.index.values

        if self.verbose is True:
            print("\nThe following {} U-axis taps were skipped due to a (very) low number of counts: ".format(len(skiptaps_u)))
            for skipped_tap in skiptaps_u:
                tapnum, counts = skipped_tap
                print("Skipped U-axis Tap {}, which had {} count(s)".format(tapnum, counts))
            print(colorama.Fore.MAGENTA + "\n... doing the same for the V axis taps {} through {}".format(taprange_v[0] + 1, taprange_v[-1] + 1))

        for tap in progressbar(taprange_v, disable=progressbar_disable, ascii=False):
            # Now do the V axis:
            tapmask_v = data[data['crsv'] == tap].index.values
            if len(tapmask_v) < 20:
                skiptaps_v.append((tap + 1, len(tapmask_v)))
                continue
            keep_v = np.isfinite(data['fb_v'][tapmask_v])

            hist_v, xbounds_v, ybounds_v = np.histogram2d(
                data['fb_v'][tapmask_v][keep_v], data['fp_v'][tapmask_v][keep_v], bins=bins)
            thresh_hist_v = self.threshold(
                hist_v, bins=bins, softening=softening)

            posx_v = np.digitize(data['fb_v'][tapmask_v], xbounds_v)
            posy_v = np.digitize(data['fp_v'][tapmask_v], ybounds_v)
            hist_mask_v = (posx_v > 0) & (posx_v <= bins[0]) & (
                posy_v > -1) & (posy_v <= bins[1])

            # Values of the histogram where the points are
            hhsub_v = thresh_hist_v[posx_v[hist_mask_v] -
                                    1, posy_v[hist_mask_v] - 1]
            pass_fb_v = data['fb_v'][tapmask_v][hist_mask_v][np.isfinite(
                hhsub_v)]

            v_axis_survivals["V Axis Tap {:02d}".format(
                tap)] = pass_fb_v.index.values

        if self.verbose is True:
            print("\nThe following {} V-axis taps were skipped due to a (very) low number of counts: ".format(len(skiptaps_v)))
            for skipped_tap in skiptaps_v:
                tapnum, counts = skipped_tap
                print("Skipped V-axis Tap {}, which had {} count(s)".format(tapnum, counts))

        # Done looping over taps

        if self.verbose is True:
            print(colorama.Fore.BLUE + "\nCollecting events that pass both U- and V-axis HyperScreen tests...", end=" ")

        u_all_survivals = np.concatenate(
            [x for x in u_axis_survivals.values()])
        v_all_survivals = np.concatenate(
            [x for x in v_axis_survivals.values()])

        # If the event passes both U- and V-axis tests, it survives
        all_survivals = np.intersect1d(u_all_survivals, v_all_survivals)
        survival_mask = np.isin(self.data.index.values, all_survivals)
        failure_mask = np.logical_not(survival_mask)

        num_survivals = sum(survival_mask)
        num_failures = sum(failure_mask)

        percent_hyperscreen_rejected = round(
            ((num_failures / self.numevents) * 100), 2)

        # Do a sanity check to look for lost events. Shouldn't be any.
        if num_survivals + num_failures != self.numevents:
            print("WARNING: Total Number of survivals and failures does \
            not equal total events in the EVT1 file. Something is wrong!")

        legacy_hyperbola_test_failures = sum(
            self.data['Hyperbola test failed'])
        percent_legacy_hyperbola_test_rejected = round(
            ((legacy_hyperbola_test_failures / self.numevents) * 100), 2)

        percent_improvement_over_legacy_test = round(
            (percent_hyperscreen_rejected - percent_legacy_hyperbola_test_rejected), 2)

        if self.verbose is True:
            print("Done")
            print(colorama.Fore.GREEN + "HyperScreen rejected" + colorama.Fore.YELLOW + " {}% of all events ({:,} bad events / {:,} total events)".format(percent_hyperscreen_rejected, sum(failure_mask), self.numevents) + colorama.Fore.GREEN +
                  "\nThe Murray+ algorithm rejects" + colorama.Fore.MAGENTA + " {}% of all events ({:,} bad events / {:,} total events)".format(percent_legacy_hyperbola_test_rejected, legacy_hyperbola_test_failures, self.numevents))

            print(colorama.Fore.GREEN + "As long as the results pass sanity checks, this is a POTENTIAL improvement of \n" +
                  colorama.Fore.BLUE + "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ POTENTIAL Improvement ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n" +
                  colorama.Fore.WHITE + "                                      {}%\n".format(percent_improvement_over_legacy_test) +
                  colorama.Fore.BLUE + "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n")

        hyperscreen_results_dict = {"ObsID": self.obsid,
                                    "Target": self.target,
                                    "Exposure Time": self.exptime,
                                    "Detector": self.detector,
                                    "Number of Events": self.numevents,
                                    "Number of Good Time Events": self.goodtimeevents,
                                    "U Axis Survivals by Tap": u_axis_survivals,
                                    "V Axis Survivals by Tap": v_axis_survivals,
                                    "U Axis All Survivals": u_all_survivals,
                                    "V Axis All Survivals": v_all_survivals,
                                    "All Survivals (event indices)": all_survivals,
                                    "All Survivals (boolean mask)": survival_mask,
                                    "All Failures (boolean mask)": failure_mask,
                                    "Percent rejected by Tapscreen": percent_hyperscreen_rejected,
                                    "Percent rejected by Hyperbola": percent_legacy_hyperbola_test_rejected,
                                    "Percent improvement": percent_improvement_over_legacy_test
                                    }

        return hyperscreen_results_dict

    def hyperbola(self, fb, a, b, h):
        """Given the normalized central tap amplitude, a, b, and h,
        return an array of length len(fb) that gives a hyperbola.

        Arguments:
            fb {[type]} -- [description]
            a {[type]} -- [description]
            b {[type]} -- [description]
            h {[type]} -- [description]

        Returns:
            [type] -- [description]
        """

        hyperbola = b * np.sqrt(((fb - h)**2 / a**2) - 1)

        return hyperbola

    def legacy_hyperbola_test(self, tolerance=0.035):
        """[summary]

        Keyword Arguments:
            tolerance {float} -- [description] (default: {0.035})

        Returns:
            [type] -- [description]
        """

        # Remind the user what tolerance they're using
        # print("{0: <25}| Using tolerance = {1}".format(" ", tolerance))

        # Set hyperbolic coefficients, depending on whether this is HRC-I or -S
        if self.detector == "HRC-I":
            a_u = 0.3110
            b_u = 0.3030
            h_u = 1.0580

            a_v = 0.3050
            b_v = 0.2730
            h_v = 1.1
            # print("{0: <25}| Using HRC-I hyperbolic coefficients: ".format(" "))
            # print("{0: <25}|    Au={1}, Bu={2}, Hu={3}".format(" ", a_u, b_u, h_u))
            # print("{0: <25}|    Av={1}, Bv={2}, Hv={3}".format(" ", a_v, b_v, h_v))

        if self.detector == "HRC-S":
            a_u = 0.2706
            b_u = 0.2620
            h_u = 1.0180

            a_v = 0.2706
            b_v = 0.2480
            h_v = 1.0710
            # print("{0: <25}| Using HRC-S hyperbolic coefficients: ".format(" "))
            # print("{0: <25}|    Au={1}, Bu={2}, Hu={3}".format(" ", a_u, b_u, h_u))
            # print("{0: <25}|    Av={1}, Bv={2}, Hv={3}".format(" ", a_v, b_v, h_v))

        # Set the tolerance boundary ("width" of the hyperbolic region)

        h_u_lowerbound = h_u * (1 + tolerance)
        h_u_upperbound = h_u * (1 - tolerance)

        h_v_lowerbound = h_v * (1 + tolerance)
        h_v_upperbound = h_v * (1 - tolerance)

        # Compute the Hyperbolae
        with np.errstate(invalid='ignore'):
            zone_u_fit = self.hyperbola(self.data["fb_u"], a_u, b_u, h_u)
            zone_u_lowerbound = self.hyperbola(
                self.data["fb_u"], a_u, b_u, h_u_lowerbound)
            zone_u_upperbound = self.hyperbola(
                self.data["fb_u"], a_u, b_u, h_u_upperbound)

            zone_v_fit = self.hyperbola(self.data["fb_v"], a_v, b_v, h_v)
            zone_v_lowerbound = self.hyperbola(
                self.data["fb_v"], a_v, b_v, h_v_lowerbound)
            zone_v_upperbound = self.hyperbola(
                self.data["fb_v"], a_v, b_v, h_v_upperbound)

        zone_u = [zone_u_lowerbound, zone_u_upperbound]
        zone_v = [zone_v_lowerbound, zone_v_upperbound]

        # Apply the masks
        # print("{0: <25}| Hyperbolic masks for U and V axes computed".format(""))

        with np.errstate(invalid='ignore'):
            # print("{0: <25}| Creating U-axis mask".format(""), end=" |")
            between_u = np.logical_not(np.logical_and(
                self.data["fp_u"] < zone_u[1], self.data["fp_u"] > -1 * zone_u[1]))
            not_beyond_u = np.logical_and(
                self.data["fp_u"] < zone_u[0], self.data["fp_u"] > -1 * zone_u[0])
            condition_u_final = np.logical_and(between_u, not_beyond_u)

            # print(" Creating V-axis mask")
            between_v = np.logical_not(np.logical_and(
                self.data["fp_v"] < zone_v[1], self.data["fp_v"] > -1 * zone_v[1]))
            not_beyond_v = np.logical_and(
                self.data["fp_v"] < zone_v[0], self.data["fp_v"] > -1 * zone_v[0])
            condition_v_final = np.logical_and(between_v, not_beyond_v)

        mask_u = condition_u_final
        mask_v = condition_v_final

        hyperzones = {"zone_u_fit": zone_u_fit,
                      "zone_u_lowerbound": zone_u_lowerbound,
                      "zone_u_upperbound": zone_u_upperbound,
                      "zone_v_fit": zone_v_fit,
                      "zone_v_lowerbound": zone_v_lowerbound,
                      "zone_v_upperbound": zone_v_upperbound}

        hypermasks = {"mask_u": mask_u, "mask_v": mask_v}

        # print("{0: <25}| Hyperbolic masks created".format(""))
        # print("{0: <25}| ".format(""))
        return hyperzones, hypermasks

    def boomerang(self, mask=None, show=True, plot_legacy_zone=True, title=None, cmap=None, savepath=None, create_subplot=False, ax=None, rasterized=True):
        """[summary]
        """
        # You can plot the image on axes of a subplot by passing
        # that axis to this function. Here are some switches to enable that.

        if create_subplot is False:
            self.fig, self.ax = plt.subplots(figsize=(12, 8))
        elif create_subplot is True:
            if ax is None:
                self.ax = plt.gca()
            else:
                self.ax = ax

        if cmap is None:
            cmap = 'plasma'

        if mask is not None:
            self.ax.scatter(self.data['fb_u'], self.data['fp_u'],
                            c=self.data['sumamps'], cmap='bone', s=0.3, alpha=0.8, rasterized=rasterized, label='All Events')

            frame = self.ax.scatter(self.data['fb_u'][mask], self.data['fp_u'][mask],
                                    c=self.data['sumamps'][mask], cmap=cmap, s=0.5, rasterized=rasterized, label='HyperScreen Accepted Events')

        else:
            frame = self.ax.scatter(self.data['fb_u'], self.data['fp_u'],
                                    c=self.data['sumamps'], cmap=cmap, s=0.5, rasterized=rasterized)

        if plot_legacy_zone is True:
            hyperzones, hypermasks = self.legacy_hyperbola_test(
                tolerance=0.035)
            self.ax.plot(self.data["fb_u"], hyperzones["zone_u_lowerbound"],
                         'o', markersize=0.3, color='black', alpha=0.8, rasterized=rasterized, label='Murray Exclusion Hyperbola')
            self.ax.plot(self.data["fb_u"], -1 * hyperzones["zone_u_lowerbound"],
                         'o', markersize=0.3, color='black', alpha=0.8, rasterized=rasterized)

            self.ax.plot(self.data["fb_u"], hyperzones["zone_u_upperbound"],
                         'o', markersize=0.3, color='black', alpha=0.8, rasterized=rasterized)
            self.ax.plot(self.data["fb_u"], -1 * hyperzones["zone_u_upperbound"],
                         'o', markersize=0.3, color='black', alpha=0.8, rasterized=rasterized)

        self.ax.grid(False)

        if title is None:
            self.ax.set_title('{} | {} | ObsID {} | {} ksec | {:,} counts'.format(
                self.target, self.detector, self.obsid, round(self.exptime / 1000, 1), self.numevents))
        else:
            self.ax.set_title(title)

        self.ax.set_ylim(-1.1, 1.1)
        self.ax.set_xlim(-0.1, 1.1)

        self.ax.set_ylabel(r'Fine Position $f_p$ $(C-A)/(A + B + C)$')
        self.ax.set_xlabel(
            r'Normalized Central Tap Amplitude $f_b$ $B / (A+B+C)$')

        if create_subplot is False:
            self.cbar = plt.colorbar(frame, pad=-0.005)
            self.cbar.set_label("SUMAMPS")

        if show is True:
            plt.show()

        if savepath is not None:
            plt.savefig(savepath, dpi=150, bbox_inches='tight')

        plt.close()

    def image(self, masked_x=None, masked_y=None, xlim=None, ylim=None, detcoords=False, title=None, cmap=None, show=True, rasterized=True, savepath=None, create_subplot=False, ax=None, nbins=(400, 400)):
        """Create a quicklook image, in detector or sky coordinates, of the
        observation. The image will be binned to 400x400 by default.

        Keyword Arguments:
            masked_x {[type]} -- [description] (default: {None})
            masked_y {[type]} -- [description] (default: {None})
            xlim {[type]} -- [description] (default: {None})
            ylim {[type]} -- [description] (default: {None})
            detcoords {bool} -- [description] (default: {False})
            title {[type]} -- [description] (default: {None})
            cmap {[type]} -- [description] (default: {None})
            show {bool} -- [description] (default: {True})
            rasterized {bool} -- [description] (default: {True})
            savepath {[type]} -- [description] (default: {None})
            create_subplot {bool} -- [description] (default: {False})
            ax {[type]} -- [description] (default: {None})
        """

        if masked_x is not None and masked_y is not None:
            x = masked_x
            y = masked_y
            img_data, yedges, xedges = np.histogram2d(y, x, nbins)
        else:
            if detcoords is False:
                x = self.data['x'][self.gtimask]
                y = self.data['y'][self.gtimask]
            elif detcoords is True:
                x = self.data['detx'][self.gtimask]
                y = self.data['dety'][self.gtimask]
            img_data, yedges, xedges = np.histogram2d(y, x, nbins)

        extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]

        # Create the Figure
        styleplots()

        # You can plot the image on axes of a subplot by passing
        # that axis to this function. Here are some switches to enable that.
        if create_subplot is False:
            self.fig, self.ax = plt.subplots()
        elif create_subplot is True:
            if ax is None:
                self.ax = plt.gca()
            else:
                self.ax = ax

        self.ax.grid(False)

        if cmap is None:
            cmap = 'viridis'

        self.ax.imshow(img_data, extent=extent, norm=LogNorm(),
                       interpolation=None, rasterized=rasterized, cmap=cmap, origin='lower')

        if title is None:
            self.ax.set_title("ObsID {} | {} | {} | {:,} events".format(
                self.obsid, self.target, self.detector, self.goodtimeevents))
        else:
            self.ax.set_title("{}".format(title))
        if detcoords is False:
            self.ax.set_xlabel("Sky X")
            self.ax.set_ylabel("Sky Y")
        elif detcoords is True:
            self.ax.set_xlabel("Detector X")
            self.ax.set_ylabel("Detector Y")

        if xlim is not None:
            self.ax.set_xlim(xlim)
        if ylim is not None:
            self.ax.set_ylim(ylim)

        if show is True:
            plt.show(block=True)

        if savepath is not None:
            plt.savefig('{}'.format(savepath))
            print("Saved image to {}".format(savepath))

        plt.close()

def styleplots():  # pragma: no cover
    """Make the plots pretty.
    """

    mpl.rcParams['agg.path.chunksize'] = 10000

    # Make things pretty
    plt.style.use('ggplot')

    labelsizes = 10

    plt.rcParams['font.size'] = labelsizes
    plt.rcParams['axes.titlesize'] = 12
    plt.rcParams['axes.labelsize'] = labelsizes
    plt.rcParams['xtick.labelsize'] = labelsizes
    plt.rcParams['ytick.labelsize'] = labelsizes