Refactor

@shevaun what do you think?