import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin


class Weight(BaseEstimator, TransformerMixin):
    
    def __init__(self, weights):
        if not isinstance(weights, np.ndarray):
            if isinstance(weights, list):
                weights = np.array(weights)
            else:
                raise ValueError("Expected list or np.ndarray")
        self.weights = weights
    
    def fit(self, X, y=None):
        return self
        
    def transform(self, X):
        if not isinstance(X, np.ndarray):
            X = np.array(X)
        if X.shape[1] != len(self.weights):
            raise ValueError("wrong feature vector length, expected " + \
                              str(len(self.weights)) + ", but actual " + \
                              str(X.shape[1]))
        return X * self.weights
    
    def fit_transform(self, X, y=None):
        return self.transform(X)
