Updated program to take argument from command line