UI updates & bug fixes