A few more bug fixes on UI