Add final code for attempt 4