Code cleanup and docs