Ui selection code improved