Updated code for displaying UI a little better.