Rework UI a little
More UI tweaks