Code cleanup WIP