Final changes to code