# Author: Kyle A. Beauchamp <kyleabeauchamp@gmail.com>
# Contributors: Robert McGibbon <rmcgibbo@gmail.com>,
#               Matthew Harrigan <matthew.p.harrigan@gmail.com>
# Copyright (c) 2014, Stanford University and the Authors
# All rights reserved.

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
from __future__ import print_function, division, absolute_import

from six.moves import cPickle
import numpy as np
import mdtraj as md
from sklearn.base import TransformerMixin
import sklearn.pipeline
from sklearn.externals.joblib import Parallel, delayed

from ..base import BaseEstimator

#-----------------------------------------------------------------------------
# Code
#-----------------------------------------------------------------------------


def featurize_all(filenames, featurizer, topology, chunk=1000, stride=1):
    """Load and featurize many trajectory files.

    Parameters
    ----------
    filenames : list of strings
        List of paths to MD trajectory files
    featurizer : Featurizer
        The featurizer to be invoked on each trajectory trajectory as
        it is loaded
    topology : str, Topology, Trajectory
        Topology or path to a topology file, used to load trajectories with
        MDTraj
    chunk : {int, None}
        If chunk is an int, load the trajectories up in chunks using
        md.iterload for better memory efficiency (less trajectory data needs
        to be in memory at once)
    stride : int, default=1
        Only read every stride-th frame.

    Returns
    -------
    data : np.ndarray, shape=(total_length_of_all_trajectories, n_features)
    indices : np.ndarray, shape=(total_length_of_all_trajectories)
    fns : np.ndarray shape=(total_length_of_all_trajectories)
        These three arrays all share the same indexing, such that data[i] is
        the featurized version of indices[i]-th frame in the MD trajectory
        with filename fns[i].
    """
    data = []
    indices = []
    fns = []

    for file in filenames:
        kwargs = {} if file.endswith('.h5') else {'top': topology}
        count = 0
        for t in md.iterload(file, chunk=chunk, stride=stride, **kwargs):
            x = featurizer.partial_transform(t)
            n_frames = len(x)

            data.append(x)
            indices.append(count + (stride*np.arange(n_frames)))
            fns.extend([file] * n_frames)
            count += (stride*n_frames)
    if len(data) == 0:
        raise ValueError("None!")

    return np.concatenate(data), np.concatenate(indices), np.array(fns)


def load(filename):
    """Load a featurizer from a cPickle file."""
    with open(filename, 'rb') as f:
        featurizer = cPickle.load(f)
    return featurizer


class Featurizer(BaseEstimator, TransformerMixin):
    """Base class for objects that featurize Trajectories.

    Notes
    -----
    At the bare minimum, a featurizer must implement the `partial_transform(traj)`
    member function.  A `transform(traj_list)` for featurizing multiple
    trajectories in batch will be provided.
    """

    def __init__(self):
        pass

    def featurize(self, traj):
        raise NotImplementedError('This API was removed. Use partial_transform instead')

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space.

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        pass

    def fit(self, traj_list, y=None):
        return self

    def transform(self, traj_list, y=None):
        """Featurize a several trajectories.

        Parameters
        ----------
        traj_list : list(mdtraj.Trajectory)
            Trajectories to be featurized.

        Returns
        -------
        features : list(np.ndarray), length = len(traj_list)
            The featurized trajectories.  features[i] is the featurized
            version of traj_list[i] and has shape
            (n_samples_i, n_features)
        """
        return [self.partial_transform(traj) for traj in traj_list]

    def save(self, filename):
        with open(filename, 'wb') as f:
            cPickle.dump(self, f)


class SuperposeFeaturizer(Featurizer):
    """Featurizer based on euclidian atom distances to reference structure.

    This featurizer transforms a dataset containing MD trajectories into
    a vector dataset by representing each frame in each of the MD trajectories
    by a vector containing the distances from a specified set of atoms to
    the 'reference position' of those atoms, in ``reference_traj``.

    Parameters
    ----------
    atom_indices : np.ndarray, shape=(n_atoms,), dtype=int
        The indices of the atoms to superpose and compute the distances with
    reference_traj : md.Trajectory
        The reference conformation to superpose each frame with respect to
        (only the first frame in reference_traj is used)
    superpose_atom_indices : np.ndarray, shape=(n_atoms,), dtype=int
        If not None, these atom_indices are used for the superposition
    """

    def __init__(self, atom_indices, reference_traj, superpose_atom_indices=None):
        self.atom_indices = atom_indices
        if superpose_atom_indices is None:
            self.superpose_atom_indices = atom_indices
        else:
            self.superpose_atom_indices = superpose_atom_indices
        self.reference_traj = reference_traj
        self.n_features = len(self.atom_indices)

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space via distance
        after superposition

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        traj.superpose(self.reference_traj, atom_indices=self.superpose_atom_indices)
        diff2 = (traj.xyz[:, self.atom_indices] -
                 self.reference_traj.xyz[0, self.atom_indices]) ** 2
        x = np.sqrt(np.sum(diff2, axis=2))
        return x


class AtomPairsFeaturizer(Featurizer):
    """Featurizer based on distances between specified pairs of atoms.

    This featurizer transforms a dataset containing MD trajectories into
    a vector dataset by representing each frame in each of the MD trajectories
    by a vector of the distances between the specified pairs of atoms.

    Parameters
    ----------
    pair_indices : np.ndarray, shape=(n_pairs, 2), dtype=int
        Each row gives the indices of two atoms involved in the interaction.
    periodic : bool, default=False
        If `periodic` is True and the trajectory contains unitcell
        information, we will compute distances under the minimum image
        convention.
    exponent : float
        Modify the distances by raising them to this exponent.
    """

    def __init__(self, pair_indices, periodic=False, exponent=1.):
        # TODO: We might want to implement more error checking here. Or during
        # featurize(). E.g. are the pair_indices supplied valid?
        self.pair_indices = pair_indices
        self.n_features = len(self.pair_indices)
        self.periodic = periodic
        self.exponent = exponent

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space via pairwise
        atom-atom distances

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        d = md.geometry.compute_distances(traj, self.pair_indices, periodic=self.periodic)
        return d ** self.exponent


class DihedralFeaturizer(Featurizer):
    """Featurizer based on dihedral angles.

    This featurizer transforms a dataset containing MD trajectories into
    a vector dataset by representing each frame in each of the MD trajectories
    by a vector containing one or more of the backbone or side-chain dihedral
    angles, or the sin and cosine of these angles.

    Parameters
    ----------
    types : list
        One or more of ['phi', 'psi', 'omega', 'chi1', 'chi2', 'chi3', 'chi4']
    sincos : bool
        Instead of outputting the angle, return the sine and cosine of the
        angle as separate features.
    """

    def __init__(self, types=['phi', 'psi'], sincos=True):
        if isinstance(types, str):
            types = [types]
        self.types = list(types)  # force a copy
        self.sincos = sincos

        known = {'phi', 'psi', 'omega', 'chi1', 'chi2', 'chi3', 'chi4'}
        if not set(types).issubset(known):
            raise ValueError('angles must be a subset of %s. you supplied %s' % (
                str(known), str(types)))

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space via calculation
        of dihedral (torsion) angles

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        x = []
        for a in self.types:
            func = getattr(md, 'compute_%s' % a)
            y = func(traj)[1]
            if self.sincos:
                x.extend([np.sin(y), np.cos(y)])
            else:
                x.append(y)
        return np.hstack(x)


class KappaAngleFeaturizer(Featurizer):
    """Featurizer to extract kappa angles.

    The kappa angle of residue `i` is the angle formed by the three CA atoms
    of residues `i-2`, `i` and `i+2`. This featurizer extracts the
    `n_residues - 4` kappa angles of each frame in a trajectory.

    Parameters
    ----------
    cos : bool
        Compute the cosine of the angle instead of the angle itself.
    """
    def __init__(self, cos=True):
        self.cos = cos

    def partial_transform(self, traj):
        ca = [a.index for a in traj.top.atoms if a.name == 'CA']
        if len(ca) < 5:
            return np.zeros((len(traj), 0), dtype=np.float32)

        angle_indices = np.array(
            [(ca[i - 2], ca[i], ca[i + 2]) for i in range(2, len(ca) - 2)])
        result = md.compute_angles(traj, angle_indices)

        if self.cos:
            return np.cos(result)

        assert result.shape == (traj.n_frames, traj.n_residues - 4)
        return result


class SASAFeaturizer(Featurizer):
    """Featurizer based on solvent-accessible surface areas.

    Parameters
    ----------

    mode : {'atom', 'residue'}, default='residue'
        In mode == 'atom', the extracted features are the per-atom
        SASA. In mode == 'residue', this is consolidated down to
        the per-residue SASA by summing over the atoms in each
        residue.

    Other Parameters
    ----------------
    probe_radius : float
    n_sphere_points : int
        If supplied, these arguments will be passed directly to
        `mdtraj.shrake_rupley`, overriding default values.

    See Also
    --------
    mdtraj.shrake_rupley
    """
    def __init__(self, mode='residue', **kwargs):
        self.mode = mode
        self.kwargs = kwargs

    def partial_transform(self, traj):
        return md.shrake_rupley(traj, mode=self.mode, **self.kwargs)


class ContactFeaturizer(Featurizer):
    """Featurizer based on residue-residue distances

    This featurizer transforms a dataset containing MD trajectories into
    a vector dataset by representing each frame in each of the MD trajectories
    by a vector of the distances between pairs of amino-acid residues.

    The exact method for computing the the distance between two residues
    is configurable with the ``scheme`` parameter.

    Parameters
    ----------
    contacts : np.ndarray or 'all'
        array containing (0-indexed) indices of the residues to compute the
        contacts for. (e.g. np.array([[0, 10], [0, 11]]) would compute
        the contact between residue 0 and residue 10 as well as
        the contact between residue 0 and residue 11.) [NOTE: if no
        array is passed then 'all' contacts are calculated. This means
        that the result will contain all contacts between residues
        separated by at least 3 residues.]
    scheme : {'ca', 'closest', 'closest-heavy'}
        scheme to determine the distance between two residues:
            'ca' : distance between two residues is given by the distance
                between their alpha carbons
            'closest' : distance is the closest distance between any
                two atoms in the residues
            'closest-heavy' : distance is the closest distance between
                any two non-hydrogen atoms in the residues
    ignore_nonprotein : bool
        When using `contact==all`, don't compute contacts between
        "residues" which are not protein (i.e. do not contain an alpha
        carbon).
    """

    def __init__(self, contacts='all', scheme='closest-heavy', ignore_nonprotein=True):
        self.contacts = contacts
        self.scheme = scheme
        self.ignore_nonprotein = ignore_nonprotein

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space via of residue-residue
        distances

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        distances, _ = md.compute_contacts(traj, self.contacts, self.scheme, self.ignore_nonprotein)
        return distances


class GaussianSolventFeaturizer(Featurizer):
    """Featurizer on weighted pairwise distance between solute and solvent.

    We apply a Gaussian kernel to each solute-solvent pairwise distance
    and sum the kernels for each solute atom, resulting in a vector
    of len(solute_indices).

    The values can be physically interpreted as the degree of solvation
    of each solute atom.

    Parameters
    ----------
    solute_indices : np.ndarray, shape=(n_solute,)
        Indices of solute atoms
    solvent_indices : np.ndarray, shape=(n_solvent,)
        Indices of solvent atoms
    sigma : float
        Sets the length scale for the gaussian kernel
    periodic : bool
        Whether to consider a periodic system in distance calculations


    References
    ----------
    ..[1] Gu, Chen, et al. BMC Bioinformatics 14, no. Suppl 2
    (January 21, 2013): S8. doi:10.1186/1471-2105-14-S2-S8.
    """

    def __init__(self, solute_indices, solvent_indices, sigma, periodic=False):
        self.solute_indices = solute_indices
        self.solvent_indices = solvent_indices
        self.sigma = sigma
        self.periodic = periodic
        self.n_features = len(self.solute_indices)

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space via calculation
        of solvent fingerprints

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        # The result vector
        fingerprints = np.zeros((traj.n_frames, self.n_features))
        atom_pairs = np.zeros((len(self.solvent_indices), 2))
        sigma = self.sigma

        for i, solute_i in enumerate(self.solute_indices):
            # For each solute atom, calculate distance to all solvent
            # molecules
            atom_pairs[:, 0] = solute_i
            atom_pairs[:, 1] = self.solvent_indices

            distances = md.compute_distances(traj, atom_pairs, periodic=True)
            distances = np.exp(-distances / (2 * sigma * sigma))

            # Sum over water atoms for all frames
            fingerprints[:, i] = np.sum(distances, axis=1)

        return fingerprints


class RawPositionsFeaturizer(Featurizer):
    """Featurize an MD trajectory into a vector space with the raw
    cartesian coordinates

    Parameters
    ----------
    atom_indices : None or array-like, dtype=int, shape=(n_atoms)
        If specified, only return the coordinates for the atoms
        given by atom_indices. Otherwise return all atoms
    ref_traj : None or md.Trajectory
        If specified, superpose each trajectory to the first frame of
        ref_traj before getting positions. If atom_indices is also
        specified, only superpose based on those atoms. The superposition
        will modify each transformed trajectory *in place*.

    """

    def __init__(self, atom_indices=None, ref_traj=None):
        super(RawPositionsFeaturizer, self).__init__()

        self.atom_indices = atom_indices

        if atom_indices is not None and ref_traj is not None:
            self.ref_traj = ref_traj.atom_slice(atom_indices)
        else:
            self.ref_traj = ref_traj


    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space with the raw
        cartesian coordinates.

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        Notes
        -----
        If you requested superposition (gave `ref_traj` in __init__) the
        input trajectory will be modified.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        # Optionally take only certain atoms
        if self.atom_indices is not None:
            p_traj = traj.atom_slice(self.atom_indices)
        else:
            p_traj = traj

        # Optionally superpose to a reference trajectory.
        if self.ref_traj is not None:
            p_traj.superpose(self.ref_traj, parallel=False)

        # Get the positions and reshape.
        value = p_traj.xyz.reshape(len(p_traj), -1)
        return value


class RMSDFeaturizer(Featurizer):
    """Featurizer based on RMSD to a series of reference frames.

    Parameters
    ----------
    trj0 : mdtraj.Trajectory
        Reference trajectory.  trj0.n_frames gives the number of features
        in this Featurizer.
    atom_indices : np.ndarray, default=None
        Which atom indices to use during RMSD calculation.  If None, MDTraj
        should default to all atoms.

    """

    def __init__(self, trj0, atom_indices=None):
        self.n_features = trj0.n_frames
        self.trj0 = trj0
        self.atom_indices = atom_indices

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space by calculating
        the RMSD to each frame in a reference trajectory.

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        X = np.zeros((traj.n_frames, self.n_features))

        for frame in range(self.n_features):
            X[:, frame] = md.rmsd(traj, self.trj0, atom_indices=self.atom_indices, frame=frame)
        return X


class DRIDFeaturizer(Featurizer):
    """Featurizer based on distribution of reciprocal interatomic
    distances (DRID)

    This featurizer transforms a dataset containing MD trajectories into
    a vector dataset by representing each frame in each of the MD trajectories
    by a vector containing the first three moments of a collection of
    reciprocal interatomic distances. For details, see [1].

    References
    ----------
    .. [1] Zhou, Caflisch; Distribution of Reciprocal of Interatomic Distances:
       A Fast Structural Metric. JCTC 2012 doi:10.1021/ct3003145

    Parameters
    ----------
    atom_indices : array-like of ints, default=None
        Which atom indices to use during DRID featurization. If None,
        all atoms are used
    """
    def __init__(self, atom_indices=None):
        self.atom_indices = atom_indices

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space using the distribution
        of reciprocal interatomic distance (DRID) method.

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """
        return md.geometry.compute_drid(traj, self.atom_indices)


class TrajFeatureUnion(BaseEstimator, sklearn.pipeline.FeatureUnion):
    """Mixtape version of sklearn.pipeline.FeatureUnion

    Notes
    -----
    Works on lists of trajectories.
    """
    def fit_transform(self, traj_list, y=None, **fit_params):
        """Fit all transformers using `trajectories`, transform the data
        and concatenate results.

        Parameters
        ----------
        traj_list : list (of mdtraj.Trajectory objects)
            Trajectories to featurize
        y : Unused
            Unused

        Returns
        -------
        Y : list (of np.ndarray)
            Y[i] is the featurized version of X[i]
            Y[i] will have shape (n_samples_i, n_features), where
            n_samples_i is the length of trajectory i and n_features
            is the total (concatenated) number of features in the
            concatenated list of featurizers.
        """
        self.fit(traj_list, y, **fit_params)
        return self.transform(traj_list)

    def transform(self, traj_list):
        """Transform traj_list separately by each transformer, concatenate results.

        Parameters
        ----------
        trajectories : list (of mdtraj.Trajectory objects)
            Trajectories to featurize

        Returns
        -------
        Y : list (of np.ndarray)
            Y[i] is the featurized version of X[i]
            Y[i] will have shape (n_samples_i, n_features), where
            n_samples_i is the length of trajectory i and n_features
            is the total (concatenated) number of features in the
            concatenated list of featurizers.

        """
        Xs = Parallel(n_jobs=self.n_jobs)(
            delayed(sklearn.pipeline._transform_one)(trans, name, traj_list, self.transformer_weights)
            for name, trans in self.transformer_list)

        X_i_stacked = [np.hstack([Xs[feature_ind][trj_ind] for feature_ind in range(len(Xs))]) for trj_ind in range(len(Xs[0]))]

        return X_i_stacked
