Kinda finalized the UI code