Refactor, implement some of @christian-marie's suggested changes