small fixes and improvement