Improved code and minor bugs fixed