Feature/ui cleanup