Code Refactor and Cleanup