Fixed shebang issue and made the script more general