Fixed bug and improved code