Improvement code and minor bug fixes