Added some changes to the controller