Polished code, hopefully fixed NPE error