# Path: msibi/utils/general.py
# def get_fn(name):
#     """Get the full path to one of the reference files shipped for testing.
# 
#     This function is taken straight from MDTraj (see https://github.com/mdtraj/mdtraj).
#     In the source distribution, these files are in ``msibi/utils/reference``,
#     but on istallation, they're moved to somewhere in the user's python
#     site-packages directory.
# 
#     Parameters
#     ----------
#     name : str
#         Name of the file ot load (with respecto to the reference/ directory).
# 
#     Examples
#     ________
#     >>> import mdtraj as md
#     >>> t = md.load(get_fun('final.hoomdxml'))
#     """
# 
#     fn = resource_filename('msibi', os.path.join('utils', 'reference', name))
# 
#     if not os.path.exists(fn):
#         raise ValueError('Sorry! %s does not exist. If you just '
#                          'added it, you\'ll have to re install' % fn)
# 
#     return fn
# 
# Path: msibi/utils/find_exclusions.py
# def find_1_n_exclusions(top, pairs, n):
#     """Find exclusions in a trajectory based on an exculsion principle
# 
#     Parameters
#     ----------
#     top : mdtraj.Topology
#         Topology object containing the types and list of bonds
#     pairs : array-like, shape=(n_pairs, 2), dtype=int
#         Each row gives the indices of two atoms.
#     n : int
#         Exclude particles in pairs separated by n or fewer bonds
#     """
#     G = nx.Graph()
#     G.add_nodes_from([a.index for a in top.atoms])
#     bonds = [b for b in top.bonds]
#     bonds_by_index = [(b[0].index, b[1].index) for b in bonds]
#     G.add_edges_from(bonds_by_index)
#     to_exclude = []
#     # TODO: make faster by looping over bonds instead of pairs
#     for i, pair in enumerate(pairs):
#         if is_1_n(pair, n, G) == True:
#             to_exclude.append(i)
#     return np.asarray(to_exclude)
# 
# Path: msibi/utils/find_exclusions.py
# def is_1_n(pair, n, G):
#     """Find if atoms in a pair are separated by n or less bonds
# 
#     Parameters
#     ----------
#     n : int
#         Return false atoms in pair are separated by n or fewer bonds
#     pair : [int, int]
#         Pair of atom indices
#     G : networkx.Graph
#         A graph with atoms and nodes and bonds as edges
# 
#     Returns
#     -------
#     answer : bool
#         answer == True if atoms are separated by n or fewer bonds, False otherwise
# 
#     The graph is expected to have atom indices as nodes, and tuples of atom indices as
#     edges. Ideally, the nodes would be MDTraj.Atom-s and edges a list of tuples of
#     MDTraj.Atom-s
#     try:
#         return n > len(nx.shortest_path(G, pair[0], pair[1])) - 1
#     except:  # no path between pair[0] and pair[1]
#         return False
#     """
#     try:
#         return n > len(nx.shortest_path(G, pair[0], pair[1])) - 1
#     except NetworkXNoPath:
#         return False


import mdtraj as md
import networkx as nx
import numpy as np
from msibi.utils.general import get_fn
from msibi.utils.find_exclusions import find_1_n_exclusions
from msibi.utils.find_exclusions import is_1_n



def test_select_pair_no_exclusions():
    """Test pair selection without exclusions"""
    top = md.load(get_fn('2chains.hoomdxml')).top
    pairs = top.select_pairs("name 'tail'", "name 'tail'")
    assert pairs.shape[0] == 190


def test_find_1_n_exclusions():
    top = md.load(get_fn('2chains.hoomdxml')).top
    pairs = top.select_pairs("name 'tail'", "name 'tail'")
    to_delete = find_1_n_exclusions(top, pairs, 3)
    assert to_delete.shape[0] == 28


def test_select_pair_with_exclusions():
    traj = md.load(get_fn('2chains.hoomdxml'))
    pairs = traj.top.select_pairs("name 'tail'", "name 'tail'")

# Path: msibi/utils/general.py
# def get_fn(name):
#     """Get the full path to one of the reference files shipped for testing.
# 
#     This function is taken straight from MDTraj (see https://github.com/mdtraj/mdtraj).
#     In the source distribution, these files are in ``msibi/utils/reference``,
#     but on istallation, they're moved to somewhere in the user's python
#     site-packages directory.
# 
#     Parameters
#     ----------
#     name : str
#         Name of the file ot load (with respecto to the reference/ directory).
# 
#     Examples
#     ________
#     >>> import mdtraj as md
#     >>> t = md.load(get_fun('final.hoomdxml'))
#     """
# 
#     fn = resource_filename('msibi', os.path.join('utils', 'reference', name))
# 
#     if not os.path.exists(fn):
#         raise ValueError('Sorry! %s does not exist. If you just '
#                          'added it, you\'ll have to re install' % fn)
# 
#     return fn
# 
# Path: msibi/utils/find_exclusions.py
# def find_1_n_exclusions(top, pairs, n):
#     """Find exclusions in a trajectory based on an exculsion principle
# 
#     Parameters
#     ----------
#     top : mdtraj.Topology
#         Topology object containing the types and list of bonds
#     pairs : array-like, shape=(n_pairs, 2), dtype=int
#         Each row gives the indices of two atoms.
#     n : int
#         Exclude particles in pairs separated by n or fewer bonds
#     """
#     G = nx.Graph()
#     G.add_nodes_from([a.index for a in top.atoms])
#     bonds = [b for b in top.bonds]
#     bonds_by_index = [(b[0].index, b[1].index) for b in bonds]
#     G.add_edges_from(bonds_by_index)
#     to_exclude = []
#     # TODO: make faster by looping over bonds instead of pairs
#     for i, pair in enumerate(pairs):
#         if is_1_n(pair, n, G) == True:
#             to_exclude.append(i)
#     return np.asarray(to_exclude)
# 
# Path: msibi/utils/find_exclusions.py
# def is_1_n(pair, n, G):
#     """Find if atoms in a pair are separated by n or less bonds
# 
#     Parameters
#     ----------
#     n : int
#         Return false atoms in pair are separated by n or fewer bonds
#     pair : [int, int]
#         Pair of atom indices
#     G : networkx.Graph
#         A graph with atoms and nodes and bonds as edges
# 
#     Returns
#     -------
#     answer : bool
#         answer == True if atoms are separated by n or fewer bonds, False otherwise
# 
#     The graph is expected to have atom indices as nodes, and tuples of atom indices as
#     edges. Ideally, the nodes would be MDTraj.Atom-s and edges a list of tuples of
#     MDTraj.Atom-s
#     try:
#         return n > len(nx.shortest_path(G, pair[0], pair[1])) - 1
#     except:  # no path between pair[0] and pair[1]
#         return False
#     """
#     try:
#         return n > len(nx.shortest_path(G, pair[0], pair[1])) - 1
#     except NetworkXNoPath:
#         return False


# Path: msibi/tests/test_select_pairs.py
import mdtraj as md
import networkx as nx
import numpy as np
from msibi.utils.general import get_fn
from msibi.utils.find_exclusions import find_1_n_exclusions
from msibi.utils.find_exclusions import is_1_n




def test_select_pair_no_exclusions():
    """Test pair selection without exclusions"""
    top = md.load(get_fn('2chains.hoomdxml')).top
    pairs = top.select_pairs("name 'tail'", "name 'tail'")
    assert pairs.shape[0] == 190


def test_find_1_n_exclusions():
    top = md.load(get_fn('2chains.hoomdxml')).top
    pairs = top.select_pairs("name 'tail'", "name 'tail'")
    to_delete = find_1_n_exclusions(top, pairs, 3)
    assert to_delete.shape[0] == 28


def test_select_pair_with_exclusions():
    traj = md.load(get_fn('2chains.hoomdxml'))
    pairs = traj.top.select_pairs("name 'tail'", "name 'tail'")

    to_delete = find_1_n_exclusions(traj.top, pairs, 3)