[Felix] Did refactoring to clean up the code.