Removed bug and Improved code