Refactor lots of code