Working code from lab