Python/Project Euler/Problem 28/sol1.py

from __future__ import print_function
from math import ceil

try:
	xrange			#Python 2
except NameError:
	xrange = range	#Python 3

def diagonal_sum(n):
	total = 1

	for i in xrange(1, int(ceil(n/2.0))):
		odd = 2*i+1
		even = 2*i
		total = total + 4*odd**2 - 6*even

	return total

if __name__ == '__main__':
	import sys

	if len(sys.argv) == 1:
		print(diagonal_sum(1001))
	else:
		try:
			n = int(sys.argv[1])
			diagonal_sum(n)
		except ValueError:
			print('Invalid entry - please enter a number')