Re-wrote script to cleanup code