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

class Decorrelator(BaseEstimator, TransformerMixin):
    def __init__(self):
        pass

    def fit(self, X):
        self.means = np.mean(X, axis=0)
        self.covmat = np.cov(X, rowvar=False, bias=True)
        self.cov_eig_vals, self.cov_eig_vecs = np.linalg.eig(self.covmat)
        return self

    def transform(self, X, zero_div_eps=1e-5):
        centered = X - self.means
        decorrelated = centered.dot(self.cov_eig_vecs)

        self.sqrt_cov_eig_vals = np.sqrt(self.cov_eig_vals + zero_div_eps)

        rescaled = decorrelated / self.sqrt_cov_eig_vals
        return rescaled

    def inverse_transform(self, X):
        if getattr(self, 'sqrt_cov_eig_vals', None) is None:
            raise Exception('First transform, then inverse')

        rescaled = X * self.sqrt_cov_eig_vals
        decorrelated = rescaled.dot(self.cov_eig_vecs.T)
        centered = decorrelated + self.means
        return centered