Refactor code and fix some ui issues