Incorporating feedback - fixed code style and comments