Small fix and added arguement option