Updated code after today's meeting