Bug Fix and improvements