Refactor/remove unused code