Minor changes to code structure and comments