# mplot1.py
# Written by Todd Wareham for CS 2500
"""
Given an upper limit nlim as a command-line argument, plots the function n*n
 on the range [0, .., (nlim-1)](increment by 1). If an additional 
 command-line argument plot-file is also given, saves plot to that file in 
 the format specified by the file's extension.

Note that this is implemented using y-data plotting, where y-data is
 explicitly created as a sequence and x-data is implicit in the indices
 of the y-data sequence.
"""

from pylab import *


if len(sys.argv) not in [2, 3]:
    print "usage: ", sys.argv[0], " nlim {plot-file}"
    sys.exit(1)

nlim = int(sys.argv[1])

## Set up plot data

y = [n*n for n in range(0, nlim)]

## Describe plot characteristics

xlabel("n")
ylabel("n*n")
title("n-squared plot")

## Create plotted line

plot(y, "b-o")

## Display plotted figure

print ">>> Creating requested plot ..."

show()

## Save plotted figure

if len(sys.argv) == 3:
    print ">>> Saving requested plot ..."
    savefig(sys.argv[2])

print ">>> Plotting finished"