wip cleaning up the code