cleaned up code based on feedback