Class fixes and cleanup