More code cleanup and updated comments