Adding UI improvements

Finish the UI prototype