bug fix  and code improvements