Code cleanup and updates