Minor refactoring in "ui.py" + update comments