Some bug fixes. Need to clean up a bit...