Big cleanup of code