Added argparse
More clean up and Unit Testing