adding code and plot for the remaining assignment