Edited controller so everything works