Cleaned up code for Code review