Refactor and code cleanup (#321)