Fixed small bug, still tweaking UI