Cleaned up code in controller.