Fixes to address bug 119

Fixes #119
Let me know if this doesn't fit the design, will re-address.