Fixing import bugs, thanks IntelliJ