UI and logic fixes