Class comments and some clean up in the code