Added some new code from teacher