
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer


def impute(data, imputerType='iterative'):
	if type(data).__name__ == 'ndarray': 
		X = data
	elif type(data).__name__ == 'wData': 
		X = data.df.values
	elif type(data).__name__ == 'DataFrame': 
		X = data.values
	else:
		raise

	imp_mean = IterativeImputer(random_state=0)
	imp_mean.fit(X)
	Xi = imp_mean.transform(X)

	if type(data).__name__ == 'ndarray': 
		return Xi
	elif type(data).__name__ == 'wData': 
		data.df[:] = Xi
		return data
	elif type(data).__name__ == 'DataFrame': 
		data[:] = Xi
		return data

