Cleaning up, fixing small issues