Bug fixing and UI changes