Bug fix and improvements