Fixes issues 6 & 8, some cleanup and further development