Improved code to run smoother