from __future__ import print_function, division, absolute_import
from mdtraj.utils.six import PY2
import abc
import re
import numpy as np
import warnings
from mdtraj import io
from msmbuilder.metrics.baseclasses import AbstractDistanceMetric, Vectorized
import inspect


class RedDimPNorm(Vectorized, AbstractDistanceMetric):

    """
    This is a class for using a reduced dimensionality representation for the trajectory data. A transformation matrix must be generated by some other method
    """

    def __init__(self, proj_obj, abs_min=None, num_vecs=None,
                 expl_var=None, which=None, metric='euclidean', p=2):
        """
        Parameters:
        -----------
        proj_obj: object
            object that has at least a method called project that has a
            call signature like:

                OBJECT.project(trajectory=None, prep_trajectory=None, which=None)

                # trajectory is an instance of msmbuilder.Trajectory
                # prep_trajectory is a prepared msmbuilder.Trajectory (np.ndarray)
                # which is a np.ndarray of which eigenvectors to use

            additionally, there needs to be an attribute called 'vals' that
            holds the eigenvalues used to determine which vectors to project
            onto
        abs_min: float, optional
            keep all eigenvectors with eigenvalues greater than <abs_min>
        num_vecs: int, optional
            number of vectors to use (sorted by decreasing eigenvalue)
        expl_var: float, optional
            keep vectors (from highest eigenvalue) such that the sum of the
            kept eigenvalues divided by the sum of all the eigenvalues is
            greater than <expl_var>
        which: np.ndarray, optional
            indices of eigenvectors to use (sorted by decreasing eigenvalue)
        metric: str, optional
            metric to use in reduced space. See Vectorized for choices
            (default: euclidean)
        p : int, optional
            Exponent for the p-norm (default: 2)


        Remarks:
        --------
        Only one of abs_min, num_vecs, expl_var, and which is necessary. If more
            than one is provided, then the one that will be used is which ever
            is given in this order: (which, num_vecs, abs_min, expl_var)
        """

        self.proj_obj = proj_obj

        if not hasattr(self.proj_obj, 'project'):
            raise Exception("proj_obj must have a method called 'project'")

        project_args = inspect.getargspec(self.proj_obj.project)

        if not ('trajectory' in project_args.args or
                'prep_trajectory' in project_args.args or
                'which' in project_args.args):

            raise Exception("proj_obj.project must take all of trajectory, "
                            "prep_trajectory, and which")

        self._set_which(abs_min=abs_min, num_vecs=num_vecs, expl_var=expl_var,
                        which=which)

        s = super(RedDimPNorm, self) if PY2 else super()
        s.__init__(metric, p)

    def _set_which(self, abs_min=None, num_vecs=None, expl_var=None, which=None):
        """
        set which eigenvectors will be selected.
        """

        if not which is None:
            self.which = np.array(which).flatten().astype(int)

        elif not num_vecs is None:
            self.which = np.arange(num_vecs)

        elif not abs_min is None:
            self.which = np.where(self.proj_obj.vals > abs_min)[0]

        elif not expl_var is None:
            tot_sum = self.proj_obj.vals.sum()
            if tot_sum <= 1e-8:
                raise Exception(
                    "total sum of eigenvalues is almost zero (or negative): %f" % tot_sum)

            self.which = np.where(np.cumsum(self.proj_obj.vals) / tot_sum < expl_var)[0]

        else:
            raise Exception("must pass one of (abs_min, num_vecs, which, or expl_var)")

    def prepare_trajectory(self, trajectory):
        """
        prepare_trajectory
        """
        proj_trajectory = self.proj_obj.project(trajectory=trajectory, which=self.which)

        return proj_trajectory
