edited controller to use less code