from euclid.metrics import Vectorized

class DWMetric(Vectorized):
    def __init__(self, dimension):
        self.metric = 'euclidean'
        self.p = None
        self._dimension = dimension
    
    def prepare_trajectory(self, trajectory):
        '''Prepare trajectory for distance computations on the *dimension*th
        diension'''
        if not self._dimension in range(0, trajectory['XYZList'].shape[1]):
            raise ValueError('Dimension mismatch')
        return trajectory['XYZList'][:, [self._dimension]]