Simplify logic
Merge pull request #2 from ahon1/ref

Simplify logic