Basic code clean up and UI fixes