Code cleanup

Plus fix sample