Major cleanup of code