UI Improvement. Moved some code around to achieve the desired UI