"""Tools for generating E3FP fingerprints.

Author: Seth Axen
E-mail: seth.axen@gmail.com
"""
from __future__ import division, print_function
import os
import logging

import numpy as np
from rdkit import Chem
import mmh3
from python_utilities.io_tools import touch_dir
from python_utilities.scripting import setup_logging
from ..config.params import get_default_value
from .fprint import Fingerprint, CountFingerprint
from .structs import Shell, shell_to_pdb
from . import array_ops


BITS = 2**32
LEVEL_DEF = get_default_value("fingerprinting", "level", int)
RADIUS_MULTIPLIER_DEF = get_default_value("fingerprinting",
                                          "radius_multiplier", float)
COUNTS_DEF = get_default_value("fingerprinting", "counts", bool)
STEREO_DEF = get_default_value("fingerprinting", "stereo", bool)
INCLUDE_DISCONNECTED_DEF = get_default_value("fingerprinting",
                                             "include_disconnected", bool)
EXCLUDE_FLOATING_DEF = get_default_value("fingerprinting",
                                         "exclude_floating", bool)
IDENT_DTYPE = np.int64  # np.dtype to use for identifiers
Y_AXIS_PRECISION = 0.1  # angstroms
Z_AXIS_PRECISION = 0.01  # rad
POLAR_CONE_RAD = np.pi / 36  # rad
MMH3_SEED = 0
BOND_TYPES = {None: 5, Chem.BondType.SINGLE: 1, Chem.BondType.DOUBLE: 2,
              Chem.BondType.TRIPLE: 3, Chem.BondType.AROMATIC: 4}

setup_logging(reset=False)


class Fingerprinter(object):

    """E3FP fingerprint generator.

    Parameters
    ----------
    bits : int or None, optional
        Maximum number of bits to which to fold returned fingerprint. Multiple
        of 2 is strongly recommended.
    level : int or None, optional
        Maximum number of iterations for fingerprint generation. If None or
        -1, run until no new substructures are identified. Because this could
        produce a different final level number for each conformer, it is
        recommended to manually specify a level.
    radius_multiplier : float, optional
        Multiple by which to increase shell size. At iteration 0, shell radius
        is 0*`radius_multiplier`, at iteration 2, radius is
        2*`radius_multiplier`, etc.
    counts : bool, optional
        Instead of simple bit-based ``Fingerprint`` object, generate
        ``CountFingerprint`` that tracks number of times each bit appears in a
        fingerprint.
    stereo : bool, optional
        Differentiate based on stereography. Resulting fingerprints are not
        comparable to non-stereo fingerprints.
    remove_duplicate_substructs : bool, optional
        If a substructure arises that corresponds to an identifier already in
        the fingerprint, then the identifier for the duplicate substructure is
        not added to fingerprint.
    include_disconnected : bool, optional:
        Include disconnected atoms from hashes and substructure. E3FP's
        advantage over ECFP relies on disconnected atoms, so the option to
        turn this off is present only for testing/comparison.
    exclude_floating : bool, optional:
        Mask atoms with no bonds (usually floating ions) from the fingerprint.
        These are often placed arbitrarily and can confound the fingerprint.

    Attributes
    ----------
    current_level : int
        The maximum level/iteration to which the fingerprinter has been run
        on the current conformer.
    level_shells : dict
        Dict matching level to set of all shells accepted at that level.
    """

    def __init__(self, bits=BITS, level=LEVEL_DEF,
                 radius_multiplier=RADIUS_MULTIPLIER_DEF, stereo=STEREO_DEF,
                 counts=COUNTS_DEF,
                 include_disconnected=INCLUDE_DISCONNECTED_DEF,
                 exclude_floating=EXCLUDE_FLOATING_DEF,
                 remove_duplicate_substructs=True):
        """Initialize fingerprinter settings."""
        self.mol = None
        if level is None:
            level = -1
        self.level = level
        if not np.log2(bits).is_integer():
            logging.warning("bits are not multiple of 2. Multiples of 2 are "
                            "highly recommended")
        self.bits = bits
        self.radius_multiplier = radius_multiplier
        if counts:
            self.fp_type = CountFingerprint
        else:
            self.fp_type = Fingerprint
        self.stereo = stereo
        self.remove_duplicate_substructs = remove_duplicate_substructs

        if self.level == -1 and not self.remove_duplicate_substructs:
            raise Exception(
                "No termination condition specified. 'level' must be "
                "provided or 'remove_duplicate_substructs' must be True")

        self.include_disconnected = include_disconnected
        self.exclude_floating = exclude_floating

        self.bond_types = BOND_TYPES
        self.reset()

    def run(self, conf=None, mol=None, return_substruct=False):
        """Generate fingerprint from provided conformer or mol and conf id.

        Parameters
        ----------
        conf : RDKit Conformer or int, optional
            Input conformer or conformer in `mol`.
        mol : RDKit Mol, optional
            Input molecule object, with at least one conformer. If `conf` not
            specified, first conformer is used.
        return_substruct : bool, optional
            Return dict mapping substructure to fingerprint indices. Keys are
            indices, values are list of substructures, represented as a tuple
            of atom indices where the first index is the central atom and the
            remaining indices (within the sphere) are sorted.
        """
        if mol is None:  # mol not provided; get from conf
            try:
                mol = conf.GetOwningMol()
            except AttributeError:  # conf is int ID; use existing mol
                mol = self.mol
        else:
            if not isinstance(conf, Chem.Conformer):
                try:
                    conf = mol.GetConformer(conf)
                except TypeError:  # conf isn't ID either. Fall back to first
                    conf = mol.GetConformer(0)

        if mol is not self.mol:
            self.reset_mol()
            self.initialize_mol(mol)
        elif conf is not self.conf:
            self.reset_conf()

        self.initialize_conformer(conf)

        for i in iter(self):
            pass

    def reset(self):
        """Clear all variables associated with the last run."""
        self.reset_mol()

    def reset_mol(self):
        """Clear all variables associated with the molecule."""
        self.atoms = None
        self.bound_atoms_dict = {}
        self.connectivity = {}
        self.init_identifiers = {}
        self.reset_conf()

    def reset_conf(self):
        """Clear only conformer-specific variables."""
        self.all_shells = []
        self.atom_coords = None
        # self.current_level = None
        self.identifiers_to_shells = {}
        self.level_shells = {}
        self.past_substructs = set()
        self.shells_gen = None

    def initialize_mol(self, mol):
        """Set general properties of `mol` that apply to all its conformers.

        Parameters
        ----------
        mol : RDKit Mol
            Input molecule ``Mol`` object.
        """
        self.mol = mol
        self.atoms = np.array([x.GetIdx() for x in mol.GetAtoms()
                               if x.GetAtomicNum() > 1])  # ignore hydrogens

        if self.exclude_floating and len(self.atoms) > 1:  # ignore floating atoms
            self.atoms = np.array([x.GetIdx() for x in mol.GetAtoms()
                                   if (x.GetAtomicNum() > 1 and
                                       x.GetDegree() > 0)])

        self.bound_atoms_dict = bound_atoms_from_mol(self.mol, self.atoms)
        self.connectivity = {}
        for i, atom1 in enumerate(self.atoms):
            for j, atom2 in enumerate(self.atoms):
                if i <= j:
                    break
                pair = (int(atom1), int(atom2))
                bond = self.mol.GetBondBetweenAtoms(*pair)
                if bond is not None:
                    bond = bond.GetBondType()
                self.connectivity[pair] = self.bond_types[bond]
                self.connectivity[pair[::-1]] = self.connectivity[pair]

        self.initialize_identifiers()

    def initialize_conformer(self, conf):
        """Retrieve atom coordinates and instantiate shells generator.

        Parameters
        ----------
        conf : RDKit Conformer
            Conformer to fingerprint
        """
        self.conf = conf
        self.atom_coords = coords_from_atoms(self.atoms, self.conf)
        self.shells_gen = ShellsGenerator(
            self.conf, self.atoms, radius_multiplier=self.radius_multiplier,
            atom_coords=self.atom_coords,
            include_disconnected=self.include_disconnected,
            bound_atoms_dict=self.bound_atoms_dict)

    def initialize_identifiers(self):
        """Set initial identifiers for atoms."""
        self.init_identifiers = identifiers_from_invariants(self.mol,
                                                            self.atoms)

    def __next__(self):
        """Run next iteration of fingerprinting."""
        if self.current_level is None:
            shells_dict = next(self.shells_gen)
            if self.current_level != 0:
                raise Exception("ShellsGenerator is not at level 0 at start."
                                " This should never happen.")

            for atom, shell in shells_dict.items():
                shell.identifier = self.init_identifiers[atom]
                self.identifiers_to_shells.setdefault(shell.identifier,
                                                      set()).add(shell)
                self.past_substructs.add(shell.substruct)

            level_shells = set(shells_dict.values())
        else:
            # stop if maximum level has been reached
            if self.current_level >= self.level and self.level != -1:
                logging.debug("Hit maximum level")
                raise StopIteration

            # stop if all substructs contain all atoms (there will never
            # be another new substruct), and that substruct has been seen
            if (self.remove_duplicate_substructs and
                all((len(x.substruct.atoms) == len(self.atoms))
                    for x in self.shells_gen.get_shells_at_level(
                        self.current_level).values())):
                logging.debug("Ran out of substructs")
                raise StopIteration

            shells_dict = next(self.shells_gen)

            for atom, shell in shells_dict.items():
                identifier = identifier_from_shell(shell, self.atom_coords,
                                                   self.connectivity,
                                                   self.current_level,
                                                   self.stereo)
                shell.identifier = identifier

            accepted_shells = sorted(shells_dict.values(),
                                     key=self._shell_to_tuple)

            # filter shells that correspond to already seen substructs
            if self.remove_duplicate_substructs:
                unique_substruct_shells = []
                for shell in accepted_shells:
                    if shell.substruct in self.past_substructs:
                        logging.debug(("Shell with identifier {} at level {} is not "
                                       "unique. Removing.").format(
                                       shell.identifier, self.current_level))
                        shell.is_duplicate = True
                        for x in self.past_substructs:
                            if x == shell.substruct:
                                shell.duplicate = x.shell
                                continue
                        continue
                    unique_substruct_shells.append(shell)
                    self.past_substructs.add(shell.substruct)

                accepted_shells = unique_substruct_shells

            # store shells
            for shell in accepted_shells:
                self.identifiers_to_shells.setdefault(shell.identifier,
                                                      set()).add(shell)

            level_shells = self.level_shells[self.current_level - 1].union(
                set(accepted_shells))
            if (len(level_shells) == len(
                    self.level_shells[self.current_level - 1])):
                self.shells_gen.back()
                logging.debug("No new shells added. Convergence reached.")
                raise StopIteration

        self.all_shells.extend(shells_dict.values())
        self.level_shells[self.current_level] = level_shells

    next = __next__

    @staticmethod
    def _shell_to_tuple(shell):
        return (shell.identifier, shell.center_atom)

    def get_shells_at_level(self, level=-1, exact=False, atom_mask=set()):
        """Get set of shells at the specified level.

        Parameters
        ----------
        level : int or None, optional
            Level/iteration
        exact : bool, optional
            Exact level
        atom_mask : int or set of int, optional
            Don't return shells whose substructures contain these atoms.

        Returns
        -------
        set of Shell : Shells at level
        """
        if level in (-1, None) or level not in self.level_shells:
            if exact or len(self.level_shells) == 0:
                raise IndexError("Level {!r} fingerprints have not yet been "
                                 "generated".format(level))

            true_level = self.current_level
            if level in (-1, None):
                if self.level not in (-1, None):
                    level = self.level
                else:
                    level = true_level
        else:
            true_level = level

        shells = self.level_shells[true_level]
        if len(atom_mask) > 0:
            try:
                atom_mask = set(atom_mask)
            except TypeError:
                atom_mask = {atom_mask}
            shells = {x for x in shells
                      if x.substruct.atoms.isdisjoint(atom_mask)}

        return shells

    def get_fingerprint_at_level(self, level=-1, bits=None, exact=False,
                                 atom_mask=set()):
        """Summary

        Parameters
        ----------
        level : int or None, optional
            Level/iteration
        bits : int or None, optional
            Return fingerprints folded to this number of bits. If unspecified,
            defaults to bits set when instantiated.
        exact : bool, optional
            Exact level
        atom_mask : int or set of int, optional
            Don't return shells whose substructures contain these atoms.

        Returns
        -------
        Fingerprint : Fingerprint at level
        """
        if bits in (-1, None):
            bits = self.bits

        shells = self.get_shells_at_level(level=level, exact=exact,
                                          atom_mask=atom_mask)

        identifiers = signed_to_unsigned_int(
            np.array([x.identifier for x in shells], dtype=IDENT_DTYPE))

        fprint = self.fp_type.from_indices(identifiers, level=level)

        return fprint.fold(bits)

    def substructs_to_pdb(self, level=None, bits=None, out_dir='substructs',
                          reorient=True, exact=False):
        """Save all accepted substructs from current level to PDB.

        Parameters
        ----------
        level : int or None, optional
            Level of fingerprinting/number of iterations
        bits : int or None, optional
            Folding level of identifiers
        out_dir : str, optional
            Directory to which to save PDB files.
        reorient : bool, optional
            Reorient substructure to match stereo quadrants.
        """
        shells = self.get_shells_at_level(level=level, exact=exact)

        if bits in (-1, None):
            bits = self.bits

        touch_dir(out_dir)

        out_files = []
        for shell in shells:
            identifier = signed_to_unsigned_int(shell.identifier) % bits
            out_file = os.path.join(out_dir, "{}.pdb.gz".format(identifier))
            shell_to_pdb(self.mol, shell, self.atom_coords,
                         self.bound_atoms_dict, out_file, reorient=reorient)
            out_files.append(out_file)
        return out_files

    @property
    def current_level(self):
        try:
            return self.shells_gen.level
        except AttributeError:
            return None

    def __iter__(self):
        return self


class ShellsGenerator(object):

    """Generate nested ``Shell`` objects from molecule upon request."""

    def __init__(self, conf, atoms, radius_multiplier=0.5,
                 include_disconnected=True, atom_coords=None,
                 bound_atoms_dict=None):
        """Initialize the generator.

        After initialization, the generator can be iterated to generate a
        ``dict`` matching atom ids to that atom's shell at that
        iteration/level.

        Parameters
        ----------
        conf : RDKit Conformer
            Conformer for which to generate shells.
        atoms : list of int
            Atom ids.
        radius_multiplier : float, optional
            Multiplier by which to increase shell radii.
        include_disconnected : bool, optional
            Include atoms not bonded to center atoms within shell.
        atom_coords : dict or None, optional
            Dict matching atom ids to 1-D array of coordinates. If None, will
            be generated.
        bound_atoms_dict : dict or None, optional
            Dict matching atom ids to set of bound atom ids. If None, will be
            generated.
        """
        self.atoms = tuple(atoms)
        self.radius_multiplier = radius_multiplier
        self.shells_dict = {}
        self.include_disconnected = include_disconnected
        self.level = None

        if atom_coords is None:
            atom_coords = coords_from_atoms(self.atoms, conf)
        atom_coords = [atom_coords.get(x) for x in self.atoms]
        self.distance_matrix = array_ops.make_distance_matrix(atom_coords)

        overlap_atoms = [(self.atoms[i], self.atoms[j]) for i, j in
                         zip(*np.where(self.distance_matrix <= array_ops.EPS))
                         if i < j]
        if len(overlap_atoms) > 0:
            logging.warning("Overlapping atoms {} in conformer {} of molecule"
                            " {}. Fingerprinting will continue but is less "
                            "reliable.".format(
                                ", ".join(map(repr, overlap_atoms)),
                                conf.GetId(),
                                conf.GetOwningMol().GetProp("_Name")))

        if not include_disconnected and bound_atoms_dict is None:
            bound_atoms_dict = bound_atoms_from_mol(conf.GetOwningMol(),
                                                    atoms)
        self.bound_atoms_dict = bound_atoms_dict

    def get_match_atoms(self, rad):
        """Get atoms within shell at radius `rad`.

        Parameters
        ----------
        rad : float
            Radius of shell.

        Returns
        -------
        dict : Dict matching atom id to set of ids for other atoms within
               shell
        """
        match_atoms_dict = {x: set() for x in self.atoms}
        atom_pair_indices_list = list(zip(*np.where(
            self.distance_matrix <= rad)))
        for i, j in atom_pair_indices_list:
            if i <= j:
                continue
            atom1, atom2 = (self.atoms[i], self.atoms[j])
            match_atoms_dict[atom1].add(atom2)
            match_atoms_dict[atom2].add(atom1)
        if not self.include_disconnected:
            match_atoms_dict = {k: v.intersection(self.bound_atoms_dict[k])
                                for k, v in match_atoms_dict.items()}
        return match_atoms_dict

    def __next__(self):
        """Get next iteration's ``dict`` of atom shells."""
        if self.level is None:
            self.level = 0
            self.shells_dict[self.level] = {x: Shell(x, radius=0.)
                                            for x in self.atoms}
            return self.shells_dict[self.level]

        self.level += 1
        self.shells_dict[self.level] = {}
        rad = self.level * self.radius_multiplier
        match_atoms_dict = self.get_match_atoms(rad)
        for atom in self.atoms:
            match_atoms = match_atoms_dict[atom]
            last_match_shells = [self.shells_dict[self.level - 1].get(x)
                                 for x in match_atoms]
            last_shell = self.shells_dict[self.level - 1][atom]
            shell = Shell(atom, last_match_shells, radius=rad,
                          last_shell=last_shell)
            self.shells_dict[self.level][atom] = shell
        return self.shells_dict[self.level]

    next = __next__

    def back(self):
        """Back up one iteration."""
        if self.level in {None, 0}:
            return
        del self.shells_dict[self.level]
        self.level -= 1

    def get_shells_at_level(self, level):
        """Get ``dict`` of atom shells at specified level/iteration.

        If not run to `level`, raises ``IndexError``.

        Parameters
        ----------
        level : int
            Level/iteration from which to retrieve shells ``dict``.

        Returns
        -------
        dict: Dict matching atom ids to that atom's ``Shell`` at that level.
        """
        if level not in self.shells_dict:
            raise IndexError(
                "Level {!r} shells have not been generated".format(level))
        return self.shells_dict[level]

    def __iter__(self):
        return self


# Getting atom properties
def coords_from_atoms(atoms, conf):
    """Build ``dict`` matching atom id to coordinates.

    Parameters
    ----------
    atoms : list of int
        Atom ids
    conf : RDKit Conformer
        Conformer from which to fetch coordinates

    Returns
    -------
    dict : Dict matching atom id to 1-D array of coordinates.
    """
    coordinates = [np.array(conf.GetAtomPosition(int(x)), dtype=np.float64)
                   for x in atoms]
    return dict(zip(atoms, coordinates))


def bound_atoms_from_mol(mol, atoms):
    """Build ``dict`` matching atom id to ids of bounded atoms.

    Bound atoms not in `atoms` are ignored.

    Parameters
    ----------
    mol : RDKit Mol
        Input mol
    atoms : list of int
        List of atom IDs

    Returns
    -------
    dict : Dict matching atom id to set of bound atom ids.
    """
    atoms = set(atoms)
    bound_atoms_dict = {}
    for bond in mol.GetBonds():
        atom1, atom2 = (bond.GetBeginAtomIdx(), bond.GetEndAtomIdx())
        if (atom1 not in atoms) or (atom2 not in atoms):
            continue
        bound_atoms_dict.setdefault(atom1, set()).add(atom2)
        bound_atoms_dict.setdefault(atom2, set()).add(atom1)
    return bound_atoms_dict


# Generating identifiers
def hash_int64_array(array, seed=MMH3_SEED):
    """Hash an int64 array into a 32-bit integer.

    Parameters
    ----------
    array : ndarray of int64
        Numpy array containing integers
    seed : any, optional
        Seed for MurmurHash3.

    Returns
    -------
    int : 32-bit integer
    """
    if not np.issubdtype(array.dtype, IDENT_DTYPE):
        raise TypeError(
            "Provided array has dtype {} not {}".format(array.dtype,
                                                        IDENT_DTYPE.__name__))
    # ensure all hashed integers are positive
    hashed_int = mmh3.hash(array, seed)
    return hashed_int


def signed_to_unsigned_int(a, bits=BITS):
    """Convert `int` between +/-`bits` to an int between 0 and `bits`.

    Parameters
    ----------
    a : int or ndarray of int
        Integer
    bits : int, optional
        Maximum size of int. E.g. 32-bit is 2^32.

    Returns
    -------
    int : unsigned integer
    """
    return (a + bits) % bits


def identifiers_from_invariants(mol, atoms):
    """Initialize ids according to Daylight invariants.

    Parameters
    ----------
    mol : RDKit Mol
        Input molecule
    atoms : list of int
        IDs for atoms in mol for which to generate identifiers.

    Returns
    -------
    ndarray of int64 : initial identifiers for atoms
    """
    identifiers = [hash_int64_array(
        invariants_from_atom(mol.GetAtomWithIdx(int(x)))) for x in atoms]
    atom_to_identifier_dict = dict(zip(atoms, identifiers))
    return atom_to_identifier_dict


def invariants_from_atom(atom):
    """Get seven invariants from atom.

    Invariants used are the six Daylight invariants, plus an indicator of
    whether the atom is in a ring, as detailed in [1].

    References
    ----------
    1. D Rogers, M Hahn. J. Chem. Inf. Model., 2010, 50 (5), pp 742-754
       https://doi.org/10.1021/ci100050t

    Parameters
    ----------
    atom : RDKit Atom
        Input atom

    Returns
    -------
    1-D array if int64: Array of 7 invariants
    """
    num_hs = atom.GetTotalNumHs()
    return np.array([atom.GetTotalDegree() - num_hs, atom.GetTotalValence() -
                     num_hs, atom.GetAtomicNum(), int(atom.GetMass()),
                     atom.GetFormalCharge(), num_hs, int(atom.IsInRing())],
                    dtype=IDENT_DTYPE)


def identifier_from_shell(shell, atom_coords, connectivity, level, stereo):
    """Determine new identifier for a shell at a specific level.

    Parameters
    ----------
    shell : Shell
        Shell for which to determine identifier
    atom_coords : dict
        Dict matching atom ids to coords.
    connectivity : dict
        Dict matching atom id pair tuples to their bond order (5 for unbound).
    level : int
        Level/iteration
    stereo : bool
        Add stereo indicators
    """
    header = [level, shell.last_shell.identifier]
    atom_tuples = atom_tuples_from_shell(shell, atom_coords, connectivity,
                                         stereo)
    flat_atom_tuples = [y for x in atom_tuples for y in x]
    arr = np.array(header + flat_atom_tuples, dtype=IDENT_DTYPE)
    return hash_int64_array(arr)


def atom_tuples_from_shell(shell, atom_coords, connectivity, stereo):
    """Generate sorted atom tuples for neighboring atoms.

    Parameters
    ----------
    shell : Shell
        Shell for which to build atom tuples
    atom_coords : dict
        Dict matching atom ids to coords.
    connectivity : dict
        Dict matching atom id pair tuples to their bond order (5 for unbound).
    stereo : bool
        Add stereo indicators to tuples
    """
    if len(shell.shells) == 0:
        return []

    # create list matching "bond order" to last iteration's atom identifier
    atom_tuples = [
        (connectivity[(shell.center_atom, x.center_atom)], x.identifier, x)
        for x in shell.shells]

    # add stereo indicator
    if stereo:
        atom_tuples.sort(key=_first_two)
        stereo_indicators = stereo_indicators_from_shell(shell, atom_tuples,
                                                         atom_coords)
        atom_tuples = [x[:-1] + (y,) + (x[-1],) for x, y
                       in zip(atom_tuples, stereo_indicators)]

    # final sort
    atom_tuples = [x[:-1] for x in atom_tuples]
    atom_tuples.sort()
    return atom_tuples


def _first_two(xs):
    return (xs[0], xs[1])


# Methods used for stereo indicators
def pick_y(atom_tuples, cent_coords, y_precision=Y_AXIS_PRECISION):
    """Pick a y-coordinate from atom tuples or mean coordinate.

    Parameters
    ----------
    atom_tuples : list of tuple
        Sorted list of atom tuples
    cent_coords : Nx3 array of float
        Coordinates of atoms with center atom at origin.
    y_precision : str, optional
        For mean to be chosen for y-coordinate, it must be at least this
        distance from the origin. Useful when atoms are symmetrical around
        the center atom where a slight shift in any atom results in a very
        different y.

    Returns
    -------
    1x3 array of float or None : y-coordinate
    int or None : index to y-atom, if y was chosen from the atoms.
    """
    y_inds = get_first_unique_tuple_inds(atom_tuples, 1,
                                         assume_sorted=True)
    # select y-axis vector
    if len(y_inds) > 0:  # unique atom could be found
        y = cent_coords[y_inds, :]
        return y, y_inds[0]
    elif len(atom_tuples) == 2:  # exactly two identical atoms
        y = cent_coords[0, :]
        return y, 0
    else:  # no unique atom could be found
        # assign y-axis as vector to mean coordinate
        y = np.mean(cent_coords, axis=0)
        if np.linalg.norm(y) < y_precision:
            return None, None
        else:
            return y, None


def pick_z(connectivity, identifiers, cent_coords, y, long_angle,
           z_precision=Z_AXIS_PRECISION):
    """Pick a z-coordinate orthogonal to `y`.

    Parameters
    ----------
    connectivity : dict
        Dict matching atom id pair tuples to their bond order (5 for unbound).
    identifiers : iterable of int
        Atom identifiers
    cent_coords : Nx3 array of float
        Coordinates of atoms with center atom at origin.
    y : 1x3 array of float
        y-coordinate
    long_angle : Nx1 array of float
        Absolute angle of atoms from orthogonal to `y`.
    z_precision : str, optional
        Minimum difference in `long_angle` between two potential z-atoms.
        Used as a tie breaker to prevent small shift in one atom resulting
        in very different z.

    Returns
    -------
    1x3 array of float or None : z-coordinate
    """
    angle_from_right = sorted(
        zip(np.asarray(long_angle / z_precision, dtype=np.int),
            connectivity,
            identifiers,
            range(len(identifiers))))

    z_angle_inds = get_first_unique_tuple_inds(angle_from_right, 1,
                                               assume_sorted=True)

    if len(z_angle_inds) > 0:
        z_ind = angle_from_right[z_angle_inds[0]][-1]
        z = array_ops.project_to_plane(cent_coords[z_ind, :], y)
        return z
    else:
        return None


def stereo_indicators_from_shell(shell, atom_tuples, atom_coords_dict,
                                 add_transform_to_shell=True):
    """Get ``list`` of ``int``s indicating location of atoms on unit sphere.

    Parameters
    ----------
    shell : Shell
        Shell for which to get stereo indicators.
    atom_tuples : list of tuple
        List of atom tuples.
    atom_coords_dict : dict
        Dict matching atom ids to coords.
    add_transform_to_shell : bool, optional
        Calculate transformation matrix to align coordinates to unit sphere,
        and add to shell.

    Returns
    -------
    list of int : stereo indicators for atoms in `atom_tuples`.
    """
    cent_coord = atom_coords_dict[shell.center_atom]

    y = None
    z = None
    if len(atom_tuples) > 0:
        (connectivity, identifiers, shells) = list(zip(*atom_tuples))

        stereo_indicators = np.zeros((len(atom_tuples),), dtype=IDENT_DTYPE)
        atoms = [x.center_atom for x in shells]
        mask = np.ones(len(atom_tuples), dtype=np.bool)

        cent_coords = np.array([atom_coords_dict.get(x) for x in atoms],
                               dtype=np.float64) - cent_coord
        # mask atom lying on center atom from consideration
        cent_overlap_indices = np.all(cent_coords == np.zeros(3), axis=1)
        mask[cent_overlap_indices] = False

        # pick y based on first unique atom tuple or mean
        y, y_ind = pick_y(atom_tuples, cent_coords)
        if y_ind is not None:
            mask[y_ind] = False  # use to ensure y not chosen for z

        if y is not None:  # y was picked
            # pick z based on closeness to pi/2 from y-axis
            long_angle = np.pi/2 - array_ops.calculate_angles(cent_coords, y)
            long_angle[np.fabs(long_angle) < array_ops.EPS] = 0.  # perfect right angles
            long_sign = np.asarray(np.sign(long_angle), dtype=IDENT_DTYPE)
            long_sign[long_sign == 0] = 1
            long_angle = np.fabs(long_angle)
            tmp_conn = np.array(connectivity, dtype=IDENT_DTYPE)[mask]
            tmp_ident = np.array(identifiers, dtype=IDENT_DTYPE)[mask]
            z = pick_z(tmp_conn, tmp_ident, cent_coords[mask], y,
                       long_angle[mask])

            if z is not None:  # z was picked
                quad_indicators = quad_indicators_from_coords(cent_coords, y,
                                                              y_ind, z,
                                                              long_sign)
                stereo_indicators = quad_indicators

            #  set indicators for atoms near poles to +/-1
            pole_inds = np.pi / 2 - long_angle < POLAR_CONE_RAD
            stereo_indicators[pole_inds] = long_sign[pole_inds]

        stereo_indicators[cent_overlap_indices] = 0
        stereo_indicators = stereo_indicators.flatten().tolist()
    else:
        stereo_indicators = []

    if add_transform_to_shell:
        shell.transform_matrix = array_ops.make_transform_matrix(cent_coord,
                                                                 y, z)

    return stereo_indicators


def quad_indicators_from_coords(cent_coords, y, y_ind, z, long_sign):
    """Create angle indicators for four quadrants in each hemisphere.

    Parameters
    ----------
    cent_coords : Nx3 array of float
        Array of centered coordinates.
    y : 1-D array of float
        Vector lying along y-axis.
    y_ind : int
        Index of `cent_coords` corresponding to `y`.
    z : 1-D array of float
        Vector lying along z-axis
    long_sign : Nx1 array of int
        Array of signs of vectors in `cent_coords` indicating whether they are
        above (+1) or below (-1) the xz-plane.

    Returns
    -------
    Nx1 array of int : Quadrant indicators. Clockwise from `z` around `y`,
                    indicators are 2, 3, 4, 5 for vectors above the xz-plane
                    and -2, -3, -4, -5 for vectors below the xz-plane.
    """
    atom_lats = array_ops.project_to_plane(cent_coords, y)

    with np.errstate(invalid='ignore'):  # y atom_lat should be (0, 0, 0)
        angle_from_z = array_ops.calculate_angles(atom_lats, z, y).flatten()

    if y_ind is not None:
        angle_from_z[y_ind] = 0.  # otherwise, will be nan

    # offset by pi/4 so z-axis isn't an edge case
    lat_angle = array_ops.rotate_angles(angle_from_z, np.pi / 4)

    # create quadrant indicators
    quad_indicators = 2 + np.asarray(lat_angle * 4/(2 * np.pi),
                                     dtype=IDENT_DTYPE)
    quad_indicators *= long_sign
    return quad_indicators


def get_first_unique_tuple_inds(tuples_list, num_ret, ignore=[],
                                assume_sorted=True):
    """Return indices of first `num_ret` unique tuples in a list.

    Only first 2 values of each tuple are considered.

    Parameters
    ----------
    tuples_list : list of tuple
        List of tuples. Only first two unique values are considered.
    num_ret : int
        Maximum number of first unique tuples to return.
    ignore : list, optional
        Indices for tuples not be considered as unique.
    assume_sorted : bool, optional
        If True, assume list is already sorted by tuples.

    Returns
    -------
    tuple of int : List of at most `num_ret` ints indicating index of
                   unique tuples in list.
    """
    unique = {}
    repeated = set(tuples_list[x][:2] for x in ignore)
    for i, tup in enumerate(tuples_list):
        val = tup[:2]
        if val in repeated:
            if assume_sorted and len(unique) == num_ret:
                break  # okay to stop looking
        elif val in unique:
            del unique[val]
            repeated.add(val)
        else:
            unique[val] = i

    return tuple(sorted(unique.values())[:num_ret])
