Fixes & some improvement