Fixed problems with animation and improved code