Code cleanup. Added comments