modified code to run better