Cleaned up code, added command line args.