bug fix and improvements