added last bits of code