Refactor some bad code