Add some comments, a bit of code clean up