Code refactoring & UI changes