Fixes to address bug 119 + 138

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