Further cleanup and improvement