did some code improvements
did some code improvements and cleanup