Code cleanup and minor changes