Corrected code ready for lecture