Code cleanup and organization