Some fixes and improvement