Code cleanup and more