Code cleanup + documentation