Some more changes to the UI. Need to fix messsages + some minor changes.