Refactor and clean up code.