Update:
Attempt at fixing the UI