Bug fixes

Random bug fixes + enhancements.