Code / UI refactor