Fix #12

Larger refactor may happen later