Code cleanup and documentation