UI updates/bug fixes