Adding comments and clean up