Improve main UI