updated code to include code review suggestions.