Fixed mistake in previous commit, optimized imports