Added some changes to classes