Cleanup code, comments and refactor a bit.