modified code after meeting with Tim