Issue #1 fixed. More code cleanup as well