General code clean up and enhancements