Fixed issue #2 + bug fixes