Major changes to code