Made changes to class structure