Cleanup. Added proper docs