UI change and bug fix