def _verify_closest_contact(traj):
    group1 = np.array([i for i in range(N_ATOMS//2)], dtype=np.int)
    group2 = np.array([i for i in range(N_ATOMS//2, N_ATOMS)], dtype=np.int)
    contact = find_closest_contact(traj, group1, group2)
    pairs = np.array([(i,j) for i in group1 for j in group2], dtype=np.int)
    dists = md.compute_distances(traj, pairs, True)[0]
    dists2 = md.compute_distances(traj, pairs, False)[0]
    nearest = np.argmin(dists)
    eq(float(dists[nearest]), contact[2], decimal=5)
    assert((pairs[nearest,0] == contact[0] and pairs[nearest,1] == contact[1]) or (pairs[nearest,0] == contact[1] and pairs[nearest,1] == contact[0]))

def distance_space(traj):
  """Convert a trajectory (or traj list) to distance space
     By default, this will compute ALL pair-wise distances and return
     a vector (or list of vectors if list of trajectories is provided)
  """
  if isinstance(traj, list):
    pairs = get_pairlist(traj[0])
    return [md.compute_distances(k,pairs) for k in traj]
  else:
    pairs = get_pairlist(traj)
    return md.compute_distances(traj,pairs)

def get_pca_array(list_chunks, topology):
    """
    Takes a list of mdtraj.Trajectory objects and featurize them to backbone -
    Alpha Carbons pairwise distances. Perform 2 component Incremental
    PCA on the featurized trajectory.

    Parameters
    ----------
    list_chunks: list of mdTraj.Trajectory objects
    topology: str
            Name of the Topology file

    Returns
    -------
    Y: np.array shape(frames, features)

    """
    pca = IncrementalPCA(n_components=2)
    top = md.load_prmtop(topology)
    ca_backbone = top.select("name CA")
    pairs = top.select_pairs(ca_backbone, ca_backbone)
    pair_distances = []
    for chunk in list_chunks:
        X = md.compute_distances(chunk, pairs)
        pair_distances.append(X)
    distance_array = np.concatenate(pair_distances)
    print("No. of data points: %d" % distance_array.shape[0])
    print("No. of features (pairwise distances): %d" % distance_array.shape[1])
    Y = pca.fit_transform(distance_array)
    return Y

def calc_obs(traj):
    arg_cz_id = 2442
    glu_cd_id = 862
    lys_nz_id = 634
    tyr_oh_id = 2019
    inactive = mdt.load("./topologies/inactive.pdb")
    active = mdt.load("./topologies/active.pdb")

    aloop_atoms_list = [i.index for residue in np.arange(147, 168) for i in inactive.topology.residue(residue).atoms]
    all_heavy = [i.index for i in inactive.topology.atoms if i.residue.is_protein and i.element.name != "hydrogen"]
    print("Processing %s" % traj)
    # load the trajectory
    trj = mdt.load(traj, atom_indices=np.arange(inactive.n_atoms))

    inactive_rms = mdt.rmsd(trj, inactive, atom_indices=all_heavy)
    active_rms = mdt.rmsd(trj, active, atom_indices=all_heavy)
    aloop_rms = mdt.rmsd(trj, inactive, frame=0, atom_indices=aloop_atoms_list)
    distances = mdt.compute_distances(trj, np.vstack(([arg_cz_id, glu_cd_id], [lys_nz_id, glu_cd_id])))
    return dict(
        fname=os.path.basename(traj),
        inactive_rmsd=inactive_rms,
        active_rmsd=active_rms,
        aloop_inactive_rmsd=aloop_rms,
        glu_arg=distances[:, 0],
        gly_lys=distances[:, 1],
    )

def indx_dists(trj, center_mol, cut):
    nmols = trj.xyz.shape[1]
    pairs = np.zeros((nmols, 2))
    pairs[:,0] = np.arange(nmols)
    pairs[:,1] = center_mol*np.ones(nmols)
    ds = mdtraj.compute_distances(trj, pairs)
    return [i for i in ds if i < cut]

def get_atom_distances(traj):
    """
    compute all atom distances for each frame in traj
    
    Parameters
    ----------
    traj : mdtraj.Trajectory
        trajectory to analyze

    Returns
    -------
    atom_distances : np.ndarray
        numpy array containing all atom distances
    atom_pairs : np.ndarray
        numpy array containing the atom pairs corresponding to
        each feature of atom_distances
    """

    n_atoms = traj.n_atoms

    atom_pairs = np.array([(i, j) for i in xrange(n_atoms) for j in xrange(i + 1, n_atoms)])

    atom_distances = md.compute_distances(traj, atom_pairs)

    return atom_distances, atom_pairs

def test_nearest_nbs(cut_off,frame_ind):
    nearest_nbs = []
    for this_ind in test.water_inds[0:2]:
        nbs = md.compute_neighbors(test.traj[frame_ind],
    cut_off,[this_ind],haystack_indices = test.water_inds)[0]
        print this_ind
        while len(nbs)!=3:
            if len(nbs) > 3:
                pairs = [[this_ind,this_nb] for this_nb in nbs]
                distances=md.compute_distances(test.traj[frame_ind],pairs)[0]
                print distances
                min_three = f(distances,3)
                print min_three
                print nbs
                nbs = [nbs[ii] for ii in min_three]
                
            else:
                print 'increase cut_off!'
        
        nbs.append(this_ind)
        nbs.sort()
        if nbs in nearest_nbs:
            print "not unique"
        else:
            nearest_nbs.append(nbs)
        
    return np.array(nearest_nbs)

 def struct_factor(self,Q,R,dt):
     """
     See equations (6) and (4) in ref [2]
     """
     water_pairs = np.array(list(combinations(sorted(self.water_inds),2)))
     step = int(dt/self.time_step)
     frame_steps= [step+random.randint(-int(step/2),int(step/2)) \
     for ii in range(int(self.total_time/dt))]
     
     while sum(frame_steps)>self.n_frames-1:
         frame_steps = frame_steps[:-1]
         
     # print ('the average time between each configuration is %f ps' % (np.mean(frame_steps)*self.time_step))
     
     N_QR = []
     current_frame = 0
     for this_step in frame_steps:
         current_frame += this_step
         water_dist = md.compute_distances(self.traj[current_frame],water_pairs)[0] # unit in nm
         N_QR.append(self.single_frame_N_QR(Q,R,water_dist))
     
     # equation (4) in ref [2]
     err_Sn_QR = np.std(N_QR)/len(N_QR)
     
     if Q == 0:
         # equation (7) in ref [2]
         Sn_QR = np.mean(N_QR)-4./3.*np.pi*self.rho*R**3.0
         
     else:
         # equations (4a-d) in ref [1]
         Sn_QR = np.mean(N_QR)-4./3.*np.pi*self.rho*3./Q**3.*(np.sin(Q*R)-Q*R*np.cos(Q*R))
     
     return Sn_QR, err_Sn_QR

    def calculate_debye_energy(self, traj, total=True):
        """Calculate the one-body burial potential

        Parameters
        ----------
        traj : mdtraj.Trajectory
            Trajectory to calculate energy over.
        sum : opt, bool
            If true (default) return the sum of the burial potentials. If
            false, return the burial energy of each individual residue.
        """
        debye = self.potential_forms["DEBYE"]

        bb_traj = self.backbone_mapping.map_traj(traj)

        r = md.compute_distances(bb_traj, self._debye_pairs)

        if total:
            Vdebye = np.zeros(bb_traj.n_frames, float)
        else:
            Vdebye = np.zeros((bb_traj.n_frames, self.n_debye_pairs), float)

        for i in range(self.n_debye_pairs):
            qi = self._debye_charges[i,0]
            qj = self._debye_charges[i,1]
            kij = self._debye_kij[i]
            if total:
                Vdebye += debye.V(r[:,i], qi, qj, kij)
            else:
                Vdebye[:, i] = debye.V(r[:,i], qi, qj, kij)
        return Vdebye

def get_square_distances(traj, aind=None):
    """
    The the atom distances in for a subset of a trajectory and 
    return it as a square distance matrix

    Parameters
    ----------
    traj : mdtraj.Trajectory 
        trajectory object
    aind : np.ndarray, optional
        atom indices to compute distances between

    Returns
    -------
    distances : np.ndarray, shape = [traj.n_frames, aind.shape[0], aind.shape[0]]
        distances between the atoms in aind
    """

    if aind is None:
        aind = np.arange(traj.n_atoms)

    pairs_ind = np.array([(i, j) for i in xrange(len(aind)) for j in xrange(i + 1, len(aind))])
    pairs = np.array([(aind[i], aind[j]) for i, j in pairs_ind])

    distances = md.compute_distances(traj, pairs)
    distances = md.geometry.squareform(distances, pairs_ind)

    return distances

def test_no_indices():
    for fn in ['2EQQ.pdb', '1bpi.pdb']:
        for opt in [True, False]:
            t = md.load(get_fn(fn))
            assert md.compute_distances(t, np.zeros((0,2), dtype=int), opt=opt).shape == (t.n_frames, 0)
            assert md.compute_angles(t, np.zeros((0,3), dtype=int), opt=opt).shape == (t.n_frames, 0)
            assert md.compute_dihedrals(t, np.zeros((0,4), dtype=int), opt=opt).shape == (t.n_frames, 0)

    def check_bonds(self):
        self.section("Bond Check (without PBCs)")

        radii = []
        pairs = []
        for (a, b) in self.topology.bonds:
            try:
                radsum = COVALENT_RADII[a.element.symbol] + COVALENT_RADII[b.element.symbol]
            except KeyError:
                raise NotImplementedError("I don't have radii information for all of your atoms")
            radii.append(radsum)
            pairs.append((a.index, b.index))

        radii = np.array(radii)
        pairs = np.array(pairs)

        distances = md.compute_distances(self.t, pairs, periodic=False)
        low, high = self.bond_low * radii, self.bond_high * radii
        extreme = np.logical_or(distances < low, distances > high)

        if np.any(extreme):
            frames, bonds = np.nonzero(extreme)
            frame, bond = frames[0], bonds[0]
            a1 = self.topology.atom(pairs[bond][0])
            a2 = self.topology.atom(pairs[bond][0])

            self.log("error: atoms (%s) and (%s) are bonded according to the topology " % (a1, a2))
            self.log("but they are a distance of %.3f nm apart in frame %d" % (distances[frame, bond], frame))
        else:
            self.log("All good.")

def get_states_Vij(model,bounds):
    """ Load trajectory, state indicators, and contact energy """

    traj = md.load("traj.xtc",top="Native.pdb")     ## Loading from file takes most time.
    rij = md.compute_distances(traj,model.contacts-np.ones(model.contacts.shape),periodic=False)
    Q = np.loadtxt("Q.dat") ## To Do: Generalize to different reaction coordinates

    state_indicator = np.zeros(len(Q),int)
    ## Assign every frame a state label. State indicator is integer 1-N for N states.
    for state_num in range(len(bounds)-1):
        instate = (Q > bounds[state_num]).astype(int)*(Q <= bounds[state_num+1]).astype(int)
        state_indicator[instate == 1] = state_num+1
    if any(state_indicator == 0):
        num_not_assign = sum((state_indicator == 0).astype(int))
        print "  Warning! %d frames were not assigned out of %d total frames!" % (num_not_assign,len(Q))
    ## Boolean arrays that indicate which state each frame is in.
    ## States are defined by their boundaries along coordinate Q.
    U  = ((Q > bounds[1]).astype(int)*(Q < bounds[2]).astype(int)).astype(bool)
    TS = ((Q > bounds[3]).astype(int)*(Q < bounds[4]).astype(int)).astype(bool)
    N  = ((Q > bounds[5]).astype(int)*(Q < bounds[6]).astype(int)).astype(bool)
    Nframes  = float(sum(N.astype(int)))
    Uframes  = float(sum(U.astype(int)))
    TSframes = float(sum(TS.astype(int)))

    Vij = model.calculate_contact_potential(rij)

    return traj,U,TS,N,Uframes,TSframes,Nframes,Vij

def compute_Epot_residues(trj, res1, res2, periodic=True, opt=True):
    """
    Compute the electrostatic energy of each pair of atoms in the two residues during the trajectory
    Return:
        Epot:       array that contains the energy per frame in kJ/mol
        ideal_dist: idealized distance of two charged particles with the residue charges and Epot
    """

    # no energy for interaction with itself
    if res1 == res2:
        Epot       = np.zeros(trj.n_frames)
        ideal_dist = np.zeros(trj.n_frames)
        return Epot, ideal_dist

    qq        = np.array([q1*q2 for (q1, q2) in itertools.product(res1.atom_charges, res2.atom_charges)])
    pairs     = itertools.product(res1.atom_indices, res2.atom_indices)
    distances = md.compute_distances(trj, pairs, periodic=periodic, opt=opt)

    Epot = (qq / distances).sum(1) # kJ / mol

    # Idealised distance of two charges with this energy
    qq_ideal = res1.charge * res2.charge
    if np.abs(qq_ideal) > 0:
        ideal_dist = (qq_ideal) / Epot 
    else:
        ideal_dist = np.zeros(trj.n_frames)

    # convert to gromacs units
    Epot *= ke_gu

    return Epot, ideal_dist

    def transform(self, traj):
        """
        Transform a trajectory into the OO features

        Parameters
        ----------
        traj : mdtraj.Trajectory

        Returns
        -------
        Xnew : np.ndarray
            sorted distances for each water molecule
        distances : np.ndarray
            distances between each water molecule
        """
        oxygens = np.array([i for i in xrange(traj.n_atoms) if traj.top.atom(i).element.symbol == 'O'])
        hydrogens = np.array([i for i in xrange(traj.n_atoms) if traj.top.atom(i).element.symbol == 'H'])

        pairsOH = np.array([(i, j) for i in oxygens for j in hydrogens])
        distances = md.compute_distances(traj, pairsOH)
        distances = distances.reshape((traj.n_frames, len(oxygens), len(hydrogens)))

        Xnew = copy.copy(distances)
        Xnew.sort()

        distances = get_square_distances(traj, oxygens)

        if not self.n_waters is None:
            Xnew = Xnew[:, :, :(2 * self.n_waters)]

        return Xnew, distances

 def obs_function(trajchunk):
     r = md.compute_distances(trajchunk,pairs,periodic=periodic)
     Econtact = np.zeros(trajchunk.n_frames,float)
     for i in range(pairs.shape[0]):
         pair_Vi = get_pair_potential(pair_type[i])
         Econtact += eps[i]*pair_Vi(r[:,i],*contact_params[i])
     return Econtact

 def map(self,traj):
     from model_builder.models.pairwise_potentials import get_pair_potential
     r = mdtraj.compute_distances(traj, self.pairs, periodic=self.periodic)
     Epair = np.zeros((traj.n_frames, self.dimension), float)
     for i in range(self.pairs.shape[0]):
         Vi = get_pair_potential(self.pair_type[i])
         Epair[:,0] += self.eps[i]*Vi(r[:,i],*self.pair_params[i])
     return Epair

 def map(self, traj):
     dists = mdtraj.compute_distances(
         traj, self.distance_indexes, periodic=self.periodic)
     res = np.zeros(
         (len(traj), self.distance_indexes.shape[0]), dtype=np.float32)
     I = np.argwhere(dists <= self.threshold)
     res[I[:, 0], I[:, 1]] = 1.0
     return res

def compute_residue_distance(trj, residues1, residues2, periodic=True, opt=True, verbose=False):

    distance_min     = np.zeros([len(residues1), len(residues2), trj.n_frames])
    distance_mean    = np.zeros([len(residues1), len(residues2), trj.n_frames])
    distance_sc_min  = np.zeros([len(residues1), len(residues2), trj.n_frames])
    distance_sc_mean = np.zeros([len(residues1), len(residues2), trj.n_frames])

    for i, res1 in enumerate(residues1):
        for j, res2 in enumerate(residues2):
            distance    = md.compute_distances(trj, itertools.product(res1.atom_indices,      res2.atom_indices),      periodic=periodic, opt=opt)
            distance_sc = md.compute_distances(trj, itertools.product(res1.sidechain_indices, res2.sidechain_indices), periodic=periodic, opt=opt)

            distance_min    [i,j,:] = distance.min(1)
            distance_mean   [i,j,:] = distance.mean(1)
            distance_sc_min [i,j,:] = distance_sc.min(1)
            distance_sc_mean[i,j,:] = distance_sc.mean(1) 

    return distance_min, distance_mean, distance_sc_min, distance_sc_mean

def extract_neighbors(chunk, pairs, n_nearest, n_A, n_B):
    cutoff = 0.52 # nanometers
    # Adapted from Ondrej
    distances = md.compute_distances(chunk, pairs, opt=True).reshape((-1, n_B, n_A))
    # structure is: idx_nearest[i_frame][i_B][i_neighbor]
    idx_nearest = [[np.where(distances[i_frame,iB,:] < cutoff)[0]
                    for iB in range(n_B)]
                    for i_frame in range(len(chunk))]
    return distances, idx_nearest