Code cleanup and improvements