Refactoring, clean up, comments, etc