'''
Compute the drift in trajectories under different distance metrics
using new distance_metrics.py code
'''
import numpy as np
from euclid.metrics import AbstractDistanceMetric


def drift(trajectory, tau, metric):
    '''
    Compute the drift in your desired metric between all pairs of conformations
    in the supplied trajectory which are seperated by tau frames.
    
    tau can be either a positive integer or an array of positive
    integers.
    
    If tau is an integer, the return value is a 1D array of length equal
    to the length of the trajectory minus tau containing the pairwise
    distance of all conformations in trajectory seperated by tau frames.
    
    If tau is an array of integers of length n, the return value is a 2D
    array with n rows and a number of columns equal to the length of the
    trajectory minus min(tau). The ith row of the returned array contains
    the pairwise distance of all of the conformations in the supplied trajectory
    separated by tau[i] frames. The final i entries in row i will be padded
    with -1s to ensure that the output 2D arrau is rectangular.
    
    prepare_kwargs=
    
    measure_kwargs=
    '''
    # make sure tau is a 1D numpy array of positive ints, or make it into one
    tau = __typecheck_tau(tau)
    if not isinstance(metric, AbstractDistanceMetric):
        raise TypeError('metric must be an instance of AbstractDistanceMetric. you supplied a %s' % metric)
    
    traj_length = trajectory['XYZList'].shape[0]
    ptraj = metric.prepare_trajectory(trajectory)
    distances = -1 * np.ones((len(tau), traj_length - np.min(tau)))
    
    for i in xrange(traj_length - np.min(tau)):
        comp_indices = filter(lambda elem: elem < traj_length, tau + i)
        d = metric.one_to_many(ptraj, ptraj, i, comp_indices)
        # these distances are the ith column
        distances[0:len(comp_indices), i] = d
        
    # if there was only 1 element in tau, reshape output so its 1D
    if distances.shape == (1, traj_length - np.min(tau)):
        distances = np.reshape(distances, traj_length - np.min(tau))
    return distances


def __typecheck_tau(tau):
    '''make sure tau is a 1D numpy array of positive ints, 
    or make it into one if possible unambiguously'''
    
    if isinstance(tau, int):
        if tau < 0: raise TypeError('Tau cannot be negative')
        tau = np.array([tau])
    else:
        tau = np.array(tau)
        if not len(tau.shape) == 1:
            raise TypeError('Tau must be a 1D array or an int. You supplied %s' % tau)
    
    # ensure positive
    if not np.all(tau == np.abs(tau)):
        raise TypeError('Taus must be all positive.')
    
    # ensure ints
    if not np.all(tau == np.array(tau, dtype='int')):
        raise TypeError('Taus must be all integers.')
    
    return tau