# (c) 2015-2018 Acellera Ltd http://www.acellera.com
# All Rights Reserved
# Distributed under HTMD Software License Agreement
# No redistribution in whole or part
#

from moleculekit.projections.projection import Projection
from moleculekit.projections.util import pp_calcDistances, pp_calcMinDistances
from moleculekit.util import ensurelist
import numpy as np
import logging

logger = logging.getLogger(__name__)


class MetricDistance(Projection):
    """ Creates a MetricDistance object

    Parameters
    ----------
    sel1 : str
        Atom selection string for the first set of atoms.
        See more `here <http://www.ks.uiuc.edu/Research/vmd/vmd-1.9.2/ug/node89.html>`__
    sel2 : str
        Atom selection string for the second set of atoms. If sel1 != sel2, it will calculate inter-set distances.
        If sel1 == sel2, it will calculate intra-set distances.
        See more `here <http://www.ks.uiuc.edu/Research/vmd/vmd-1.9.2/ug/node89.html>`__
    groupsel1 : ['all','residue'], optional
        Group all atoms in `sel1` to the single minimum distance. Alternatively can calculate the minimum distance of a
        residue containing the atoms in sel1.
    groupsel2 : ['all','residue'], optional
        Same as groupsel1 but for `sel2`
    metric : ['distances','contacts'], optional
        Set to 'contacts' to calculate contacts instead of distances
    threshold : float, optional
        The threshold under which a distance is considered in contact. Units in Angstrom.
    pbc : bool, optional
        Set to false to disable distance calculations using periodic distances
    truncate : float, optional
        Set all distances larger than `truncate` to `truncate`. Units in Angstrom.
    update :
        Not functional yet

    Returns
    -------
    proj : MetricDistance object
    """

    def __init__(
        self,
        sel1,
        sel2,
        groupsel1=None,
        groupsel2=None,
        metric="distances",
        threshold=8,
        pbc=True,
        truncate=None,
    ):
        super().__init__()

        self.sel1 = sel1
        self.sel2 = sel2
        self.groupsel1 = groupsel1
        self.groupsel2 = groupsel2
        self.metric = metric
        self.threshold = threshold
        self.pbc = pbc
        self.truncate = truncate

    def _checkChains(self, mol, sel1, sel2):
        if np.array_equal(sel1, sel2):
            return
        if np.ndim(sel1) == 1:
            sel1 = sel1[None, :]
        if np.ndim(sel2) == 1:
            sel2 = sel2[None, :]
        sel1 = np.any(sel1, axis=0)
        sel2 = np.any(sel2, axis=0)
        chains_sel1 = mol.chain[sel1]
        chains_sel2 = mol.chain[sel2]
        if len(np.intersect1d(chains_sel1, chains_sel2)):
            logger.warning(
                "Atomselections sel1 and sel2 of MetricDistance contain atoms belonging to a common chain. "
                "Atoms within the same chain will not have periodic distances computed. "
                "Ensure that chains are properly defined in your topology file."
            )

    def project(self, mol):
        """ Project molecule.

        Parameters
        ----------
        mol : :class:`Molecule <moleculekit.molecule.Molecule>`
            A :class:`Molecule <moleculekit.molecule.Molecule>` object to project.

        Returns
        -------
        data : np.ndarray
            An array containing the projected data.
        """
        getMolProp = lambda prop: self._getMolProp(mol, prop)
        sel1 = getMolProp("sel1")
        sel2 = getMolProp("sel2")
        if self.pbc:
            self._checkChains(mol, sel1, sel2)

        if np.ndim(sel1) == 1 and np.ndim(sel2) == 1:  # normal distances
            metric = pp_calcDistances(
                mol,
                sel1,
                sel2,
                self.metric,
                self.threshold,
                self.pbc,
                truncate=self.truncate,
            )
        else:  # minimum distances by groups
            metric = pp_calcMinDistances(
                mol,
                sel1,
                sel2,
                self.metric,
                self.threshold,
                pbc=self.pbc,
                truncate=self.truncate,
            )

        return metric

    def _calculateMolProp(self, mol, props="all"):
        props = ("sel1", "sel2") if props == "all" else props
        res = {}
        if "sel1" in props:
            res["sel1"] = self._processSelection(mol, self.sel1, self.groupsel1)
        if "sel2" in props:
            res["sel2"] = self._processSelection(mol, self.sel2, self.groupsel2)
        return res

    def _processSelection(self, mol, sel, groupsel):
        # If user passes simple string selections or 1D array of ints or bools
        if isinstance(sel, str) or (
            isinstance(sel, np.ndarray)
            and sel.ndim == 1
            and (np.issubdtype(sel.dtype, np.integer) or sel.dtype == bool)
        ):
            if groupsel is None:
                sel = mol.atomselect(sel)
            elif groupsel == "all":
                sel = self._processMultiSelections(mol, [sel])
            elif groupsel == "residue":
                sel = self._groupByResidue(mol, sel)
            else:
                raise NameError("Invalid groupsel argument")
        elif isinstance(sel, np.ndarray) or isinstance(
            sel, list
        ):  # If user passes his own sets of groups
            sel = self._processMultiSelections(mol, sel)
        else:
            raise RuntimeError(
                "Invalid atom selection. Either provide a string, a list of string, a 1D numpy array (int/bool) or a 2D numpy array for groups."
            )

        if np.sum(sel) == 0:
            raise NameError("Selection returned 0 atoms")
        return sel

    def _processMultiSelections(self, mol, sel):
        newsel = np.zeros((len(sel), mol.numAtoms), dtype=bool)
        for s in range(len(sel)):
            if isinstance(sel[s], str):
                newsel[s, :] = mol.atomselect(sel[s])
            elif isinstance(sel[s], np.ndarray) and (
                np.issubdtype(sel[s].dtype, np.integer) or sel[s].dtype == bool
            ):
                newsel[s, sel[s]] = True
            else:
                raise RuntimeError("Invalid selection provided for groups")
        return newsel

    def _groupByResidue(self, mol, sel):
        import pandas as pd

        idx = mol.atomselect(sel, indexes=True)
        df = pd.DataFrame({"a": mol.resid[idx]})
        gg = df.groupby(by=df.a).groups  # Grouping by same resids

        newsel = np.zeros((len(gg), mol.numAtoms), dtype=bool)
        for i, res in enumerate(sorted(gg)):
            newsel[
                i, idx[gg[res]]
            ] = True  # Setting the selected indexes to True which correspond to the same residue
        return newsel

    def getMapping(self, mol):
        """ Returns the description of each projected dimension.

        Parameters
        ----------
        mol : :class:`Molecule <moleculekit.molecule.Molecule>` object
            A Molecule object which will be used to calculate the descriptions of the projected dimensions.

        Returns
        -------
        map : :class:`DataFrame <pandas.core.frame.DataFrame>` object
            A DataFrame containing the descriptions of each dimension
        """
        getMolProp = lambda prop: self._getMolProp(mol, prop)
        sel1 = getMolProp("sel1")
        sel2 = getMolProp("sel2")

        if np.ndim(sel1) == 2:
            protatoms = []
            for i in range(sel1.shape[0]):
                protatoms.append(np.where(sel1[i, :] == True)[0])
        else:
            protatoms = np.where(sel1)[0]
        if np.ndim(sel2) == 2:
            ligatoms = []
            for i in range(sel2.shape[0]):
                ligatoms.append(np.where(sel2[i, :] == True)[0])
        else:
            ligatoms = np.where(sel2)[0]

        numatoms1 = len(protatoms)
        numatoms2 = len(ligatoms)

        from pandas import DataFrame

        types = []
        indexes = []
        description = []
        if np.array_equal(sel1, sel2):
            for i in range(numatoms1):
                for j in range(i + 1, numatoms1):
                    atm1 = protatoms[i]
                    atm2 = protatoms[j]
                    desc = "{} between {} {} {} and {} {} {}".format(
                        self.metric[:-1],
                        mol.resname[atm1],
                        mol.resid[atm1],
                        mol.name[atm1],
                        mol.resname[atm2],
                        mol.resid[atm2],
                        mol.name[atm2],
                    )
                    types += [self.metric[:-1]]
                    indexes += [[atm1, atm2]]
                    description += [desc]
        else:
            for i in range(numatoms1):
                for j in range(numatoms2):
                    atm1 = protatoms[i]
                    atm2 = ligatoms[j]
                    desc = "{} between {} {} {} and {} {} {}".format(
                        self.metric[:-1],
                        mol.resname[atm1],
                        mol.resid[atm1],
                        mol.name[atm1],
                        mol.resname[atm2],
                        mol.resid[atm2],
                        mol.name[atm2],
                    )
                    types += [self.metric[:-1]]
                    indexes += [[atm1, atm2]]
                    description += [desc]
        return DataFrame(
            {"type": types, "atomIndexes": indexes, "description": description}
        )


class MetricSelfDistance(MetricDistance):
    def __init__(
        self,
        sel,
        groupsel=None,
        metric="distances",
        threshold=8,
        pbc=True,
        truncate=None,
    ):
        """ Creates a MetricSelfDistance object

        Parameters
        ----------
        sel : str
            Atom selection string for which to calculate the self distance.
            See more `here <http://www.ks.uiuc.edu/Research/vmd/vmd-1.9.2/ug/node89.html>`__
        groupsel : ['all','residue'], optional
            Group all atoms in `sel` to the single minimum distance. Alternatively can calculate the minimum distance
            of a residue containing the atoms in `sel`.
        metric : ['distances','contacts'], optional
            Set to 'contacts' to calculate contacts instead of distances
        threshold : float, optional
            The threshold under which a distance is considered in contact
        pbc : bool, optional
            Set to false to disable distance calculations using periodic distances
        truncate : float, optional
            Set all distances larger than `truncate` to `truncate`
        update :
            Not functional yet

        Returns
        -------
        proj : MetricDistance object
        """
        super().__init__(sel, sel, groupsel, groupsel, metric, threshold, pbc, truncate)


def contactVecToMatrix(vector, atomIndexes):
    from copy import deepcopy

    if np.ndim(vector) != 1:
        raise RuntimeError(
            "Please pass a 1D vector to the contactVecToMatrix function."
        )
    # Calculating the unique atom groups in the mapping
    uqAtomGroups = []
    atomIndexes = deepcopy(list(atomIndexes))
    for ax in atomIndexes:
        ax[0] = ensurelist(ax[0])
        ax[1] = ensurelist(ax[1])
        if ax[0] not in uqAtomGroups:
            uqAtomGroups.append(ax[0])
        if ax[1] not in uqAtomGroups:
            uqAtomGroups.append(ax[1])
    uqAtomGroups.sort(key=lambda x: x[0])  # Sort by first atom in each atom list
    num = len(uqAtomGroups)

    matrix = np.zeros((num, num), dtype=vector.dtype)
    mapping = np.ones((num, num), dtype=int) * -1
    for i in range(len(vector)):
        row = uqAtomGroups.index(atomIndexes[i][0])
        col = uqAtomGroups.index(atomIndexes[i][1])
        matrix[row, col] = vector[i]
        matrix[col, row] = vector[i]
        mapping[row, col] = i
        mapping[col, row] = i
    return matrix, mapping, uqAtomGroups


def reconstructContactMap(
    vector,
    mapping,
    truecontacts=None,
    plot=True,
    figsize=(7, 7),
    dpi=80,
    title=None,
    outfile=None,
    colors=None,
):
    """ Plots a given vector as a contact map

    Parameters
    ----------
    vector : np.ndarray or list
        A 1D vector of contacts
    mapping : pd.DataFrame
        A pandas DataFrame which describes the dimensions of the projection
    truecontacts : np.ndarray or list
        A 1D vector of true contacts
    plot : bool
        To plot or not to plot
    figsize : tuple
        The size of the final plot in inches
    dpi : int
        Dots per inch
    outfile : str
        Path of file in which to save the plot

    Returns
    -------
    cm : np.ndarray
        The input vector converted into a 2D numpy array

    Examples
    --------
    >>> reconstructContactMap(contacts, mapping)
    To use it with distances instead of contacts pass ones as the concat vector
    >>> reconstructContactMap(np.ones(dists.shape, dtype=bool), mapping, colors=dists)
    """
    from copy import deepcopy
    from matplotlib import cm as colormaps

    if np.ndim(vector) != 1:
        raise RuntimeError(
            "Please pass a 1D vector to the reconstructContactMap function."
        )

    if truecontacts is None:
        truecontacts = np.zeros(len(vector), dtype=bool)
    if len(vector) != len(mapping):
        raise RuntimeError("Vector and map length must match.")

    # Checking if contacts or distances exist in the data
    contactidx = mapping.type == "contact"
    if not np.any(contactidx):
        contactidx = mapping.type == "distance"
        if not np.any(contactidx):
            raise RuntimeError(
                "No contacts or distances found in the MetricData object. Check the `.map` property of the object for a description of your projection."
            )

    # Creating the 2D contact maps
    cm, newmapping, uqAtomGroups = contactVecToMatrix(vector, mapping.atomIndexes)
    cmtrue, _, _ = contactVecToMatrix(truecontacts, mapping.atomIndexes)
    num = len(uqAtomGroups)

    if plot:
        from matplotlib import pylab as plt

        f = plt.figure(figsize=figsize, dpi=dpi)
        plt.imshow(
            cmtrue / 2,
            interpolation="none",
            vmin=0,
            vmax=1,
            aspect="equal",
            cmap="Greys",
        )  # /2 to convert to gray from black

        rows, cols = np.where(cm)
        colorbar = False
        if colors is None:
            truecms = vector & truecontacts
            colors = np.array(["r"] * len(vector), dtype=object)
            colors[truecms] = "#ffff00"
        elif isinstance(colors, np.ndarray) and isinstance(colors[0], float):
            mpbl = colormaps.ScalarMappable(cmap=colormaps.jet)
            mpbl.set_array(colors)
            colors = mpbl.to_rgba(colors)
            colorbar = True
        if len(colors) == len(vector):
            colors = colors[newmapping[rows, cols]]

        plt.scatter(rows, cols, s=figsize[0] * 5, marker="o", c=colors, lw=0)
        if colorbar:
            plt.colorbar(mpbl)

        ax = f.axes[0]
        # Major ticks
        ax.set_xticks(np.arange(0, num, 1))
        ax.set_yticks(np.arange(0, num, 1))

        # Labels for major ticks
        ax.set_xticklabels([x[0] for x in uqAtomGroups])
        ax.set_yticklabels([x[0] for x in uqAtomGroups],)

        # Minor ticks
        ax.set_xticks(np.arange(-0.5, num, 1), minor=True)
        ax.set_yticks(np.arange(-0.5, num, 1), minor=True)
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=90)

        # Gridlines based on minor ticks
        ax.grid(which="minor", color="#969696", linestyle="-", linewidth=1)
        ax.tick_params(axis="both", which="both", length=0)
        plt.xlim([-0.5, num - 0.5])
        plt.ylim([-0.5, num - 0.5])
        plt.xlabel("Atom index")
        plt.ylabel("Atom index")
        if title:
            plt.title(title)
        if outfile is not None:
            plt.savefig(outfile, dpi=dpi, bbox_inches="tight", pad_inches=0.2)
            plt.close()
        else:
            plt.show()
    return cm


import unittest
from moleculekit.home import home
import os


class _TestMetricDistance(unittest.TestCase):
    @classmethod
    def setUpClass(self):
        from moleculekit.molecule import Molecule
        from os import path

        self.mol = Molecule(
            path.join(home(dataDir="test-projections"), "trajectory", "filtered.pdb")
        )
        self.mol_skipped = self.mol.copy()

        self.mol.read(
            path.join(home(dataDir="test-projections"), "trajectory", "traj.xtc")
        )
        self.mol_skipped.read(
            path.join(home(dataDir="test-projections"), "trajectory", "traj.xtc"),
            skip=10,
        )

    def test_contacts(self):
        metr = MetricDistance(
            "protein and name CA", "resname MOL and noh", metric="contacts", threshold=8
        )
        data = metr.project(self.mol)
        reffile = os.path.join(
            home(dataDir="test-projections"), "metricdistance", "distances.npy"
        )
        refdata = np.load(reffile) < 8

        assert np.allclose(data, refdata, atol=1e-3), "Contact calculation is broken"

    def test_atomselects(self):
        # String atomselection
        metr = MetricSelfDistance("protein and name CA", metric="contacts", threshold=8)
        data = metr.project(self.mol_skipped)
        mapping = metr.getMapping(self.mol_skipped)
        assert data.shape == (20, 38226) and mapping.shape == (38226, 3)

        # Integer index atomselection
        ca_atoms = self.mol_skipped.atomselect("protein and name CA", indexes=True)
        metr = MetricSelfDistance(ca_atoms, metric="contacts", threshold=8)
        data = metr.project(self.mol_skipped)
        mapping = metr.getMapping(self.mol_skipped)
        assert data.shape == (20, 38226) and mapping.shape == (38226, 3)

        # Boolean atomselection
        ca_atoms = self.mol_skipped.atomselect("protein and name CA")
        metr = MetricSelfDistance(ca_atoms, metric="contacts", threshold=8)
        data = metr.project(self.mol_skipped)
        mapping = metr.getMapping(self.mol_skipped)
        assert data.shape == (20, 38226) and mapping.shape == (38226, 3)

        # Manual group atomselection with integers
        ca_atoms = self.mol_skipped.atomselect("protein and name CA", indexes=True)
        metr = MetricSelfDistance([ca_atoms[0::3], ca_atoms[1::3], ca_atoms[2::3]],)
        data = metr.project(self.mol_skipped)
        mapping = metr.getMapping(self.mol_skipped)
        assert data.shape == (20, 3) and mapping.shape == (3, 3)

        # Manual group atomselection with booleans
        ca_atoms = self.mol_skipped.atomselect("protein and name CA", indexes=True)
        ca_bools1 = self.mol_skipped.atomselect(
            "index {}".format(" ".join(map(str, ca_atoms[0::3])))
        )
        ca_bools2 = self.mol_skipped.atomselect(
            "index {}".format(" ".join(map(str, ca_atoms[1::3])))
        )
        ca_bools3 = self.mol_skipped.atomselect(
            "index {}".format(" ".join(map(str, ca_atoms[2::3])))
        )
        metr = MetricSelfDistance([ca_bools1, ca_bools2, ca_bools3])
        newdata = metr.project(self.mol_skipped)
        mapping = metr.getMapping(self.mol_skipped)
        assert (
            newdata.shape == (20, 3)
            and mapping.shape == (3, 3)
            and np.array_equal(data, newdata)
        )

    def test_distances_trivial(self):
        from moleculekit.molecule import Molecule
        import numpy as np

        mol = Molecule().empty(3)
        mol.name[:] = "C"
        mol.element[:] = "C"
        mol.chain[:] = list(
            map(str, range(3))
        )  # If they are in the same chain, no wrapping is done for distances
        mol.coords = np.zeros(
            (3, 3, 2), dtype=np.float32
        )  # Make two frames so we check if the code works for nframes
        mol.coords[1, :, 0] = [3, 3, 3]
        mol.coords[2, :, 0] = [5, 5, 5]
        mol.coords[1, :, 1] = [7, 7, 7]
        mol.coords[2, :, 1] = [6, 6, 6]

        realdistances = np.linalg.norm(mol.coords[[1, 2], :, :], axis=1).T

        metr = MetricDistance("index 0", "index 1 2", metric="distances", pbc=False)
        data = metr.project(mol)
        assert np.allclose(
            data, realdistances
        ), "Trivial distance calculation is broken"

        # Test wrapped distances
        wrappedcoords = np.mod(mol.coords, 2)
        wrappedrealdistances = np.linalg.norm(wrappedcoords[[1, 2], :, :], axis=1).T

        mol.box = np.full((3, 2), 2, dtype=np.float32)  # Make box 2x2x2A large
        metr = MetricDistance("index 0", "index 1 2", metric="distances", pbc=True)
        data = metr.project(mol)
        assert np.allclose(
            data, wrappedrealdistances
        ), "Trivial wrapped distance calculation is broken"

        # Test min distances
        metr = MetricDistance(
            "index 0",
            "index 1 2",
            metric="distances",
            pbc=False,
            groupsel1="all",
            groupsel2="all",
        )
        data = metr.project(mol)
        assert np.allclose(
            data.flatten(), np.min(realdistances, axis=1)
        ), "Trivial distance calculation is broken"

        # Test ordering
        mol = Molecule().empty(4)
        mol.name[:] = "C"
        mol.element[:] = "C"
        mol.chain[:] = list(
            map(str, range(4))
        )  # If they are in the same chain, no wrapping is done for distances
        mol.coords = np.zeros(
            (4, 3, 2), dtype=np.float32
        )  # Make two frames so we check if the code works for nframes
        mol.coords[1, :, 0] = [1, 1, 1]
        mol.coords[2, :, 0] = [3, 3, 3]
        mol.coords[3, :, 0] = [5, 5, 5]
        mol.coords[1, :, 1] = [1, 1, 1]
        mol.coords[2, :, 1] = [7, 7, 7]
        mol.coords[3, :, 1] = [6, 6, 6]

        realdistances = np.linalg.norm(
            mol.coords[[2, 3], :, :] - mol.coords[0], axis=1
        ).T
        realdistances = np.hstack(
            (
                realdistances,
                np.linalg.norm(mol.coords[[2, 3], :, :] - mol.coords[1], axis=1).T,
            )
        )

        metr = MetricDistance("index 0 1", "index 2 3", metric="distances", pbc=False)
        data = metr.project(mol)
        assert np.allclose(
            data, realdistances
        ), "Trivial distance calculation has broken ordering"

    def test_distances(self):
        metr = MetricDistance(
            "protein and name CA", "resname MOL and noh", metric="distances"
        )
        data = metr.project(self.mol)
        refdata = np.load(
            os.path.join(
                home(dataDir="test-projections"), "metricdistance", "distances.npy"
            )
        )
        assert np.allclose(data, refdata, atol=1e-3), "Distance calculation is broken"

    def test_mindistances(self):
        metr = MetricDistance(
            "protein and noh",
            "resname MOL and noh",
            groupsel1="residue",
            groupsel2="all",
        )
        data = metr.project(self.mol)
        refdata = np.load(
            os.path.join(
                home(dataDir="test-projections"), "metricdistance", "mindistances.npy"
            )
        )
        assert np.allclose(
            data, refdata, atol=1e-3
        ), "Minimum distance calculation is broken"

    def test_mindistances_truncate(self):
        metr = MetricDistance(
            "protein and noh",
            "resname MOL and noh",
            groupsel1="residue",
            groupsel2="all",
            truncate=3,
        )
        data = metr.project(self.mol)
        refdata = np.load(
            os.path.join(
                home(dataDir="test-projections"), "metricdistance", "mindistances.npy"
            )
        )
        assert np.allclose(
            data, np.clip(refdata, 0, 3), atol=1e-3
        ), "Minimum distance calculation is broken"

    def test_selfmindistance_manual(self):
        metr = MetricDistance(
            "protein and resid 1 to 50 and noh",
            "protein and resid 1 to 50 and noh",
            groupsel1="residue",
            groupsel2="residue",
        )
        data = metr.project(self.mol)
        refdata = np.load(
            os.path.join(
                home(dataDir="test-projections"),
                "metricdistance",
                "selfmindistance.npy",
            )
        )
        assert np.allclose(data, refdata, atol=1e-3), "Manual self-distance is broken"

    def test_selfmindistance_auto(self):
        metr = MetricSelfDistance(
            "protein and resid 1 to 50 and noh", groupsel="residue"
        )
        data = metr.project(self.mol)
        refdata = np.load(
            os.path.join(
                home(dataDir="test-projections"),
                "metricdistance",
                "selfmindistance.npy",
            )
        )
        assert np.allclose(
            data, refdata, atol=1e-3
        ), "Automatic self-distance is broken"

    def test_mindistances_skip(self):
        metr = MetricSelfDistance(
            "protein and resid 1 to 50 and noh", groupsel="residue"
        )
        data = metr.project(self.mol_skipped)
        refdata = np.load(
            os.path.join(
                home(dataDir="test-projections"),
                "metricdistance",
                "selfmindistance.npy",
            )
        )
        assert np.allclose(
            data, refdata[::10, :], atol=1e-3
        ), "Minimum distance calculation with skipping is broken"

    def test_reconstruct_contact_map(self):
        from moleculekit.util import tempname
        from moleculekit.molecule import Molecule
        import matplotlib

        matplotlib.use("Agg")

        mol = Molecule("1yu8")

        metr = MetricSelfDistance("protein and name CA", metric="contacts", pbc=False)
        data = metr.project(mol)[0]
        mapping = metr.getMapping(mol)

        tmpfile = tempname(suffix=".svg")
        cm, newmapping, uqAtomGroups = contactVecToMatrix(data, mapping.atomIndexes)
        reconstructContactMap(data, mapping, outfile=tmpfile)
        refdata = np.load(
            os.path.join(
                home(dataDir="test-projections"),
                "metricdistance",
                "contactvectomatrix.npz",
            )
        )

        assert np.array_equal(refdata["cm"], cm)
        assert np.array_equal(refdata["newmapping"], newmapping)
        assert np.array_equal(refdata["uqAtomGroups"], uqAtomGroups)

    def test_description(self):
        metr = MetricDistance("protein and noh", "resname MOL and noh", truncate=3)
        data = metr.project(self.mol)
        atomIndexes = metr.getMapping(self.mol).atomIndexes.values
        refdata = np.load(
            os.path.join(
                home(dataDir="test-projections"), "metricdistance", "description.npy"
            ),
            allow_pickle=True,
        )
        assert np.array_equal(refdata, atomIndexes)


if __name__ == "__main__":
    unittest.main(verbosity=2)
