Fixing imports and other issues