More correct verison of code