WIP, code cleaning