Code clean up and refactors