Cleaned up code for improved running