Some cleanup & small changes