Improve UI - closed #14