Some cleanup and fixing