#!/bin/env python3

import argparse
import matplotlib.pyplot as plt
import math

def squaresR2(x):
	count = 0;
	i = math.ceil(math.sqrt(x))
	while i > 0:
		# rozkłady z zerem zostaną policzone tylko raz
		# dzięki temu po przemnożeniu uzyskamy dobry wynik
		j = 0
		while j*j <= x:
			if j*j + i*i == x:
				count += 1
			j += 1
		i -= 1
	# każda z liczb może być ujemna
	# więc każda kombinacja ma 4 różne warianty
	return count*4

def jacobiR2(x):
	count = 0;
	for i in range(1, x+1, 2):
		if x%i != 0:
			continue
		elif (i-1) % 4 == 0:
			count += 1
		elif (i-3) % 4 == 0:
			count -= 1
	return count*4;

def plot(start=0, end=None, mid=None, delta=50, strategy="both"):
	if mid is not None:
		start = mid - delta
		end = mid + delta + 1
	if end is None or end < start:
		end = start + delta
	if strategy == "jacobi":
		R = [jacobiR2(i) for i in range(start, end)]
	else:
		R = [squaresR2(i) for i in range(start, end)]
	if strategy == "both":
		# check whether our results match up with Jacobi's Theorem
		R2 = [jacobiR2(i) for i in range(start, end)]
		confirmed = True
		for i in range(end-start):
			if R[i] != R2[i]:
				print("Our result for {} doesn't match Jacobi's Theorem".format(i))
				confirmed = False
		if confirmed:
			print("both strategies gave the same result")
	# make a plot
	plt.scatter(range(start, end), R, marker=".")
	plt.ylim(min(R)-1, max(R)+1)
	plt.show()

if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument("--start", type=int, default=0)
	parser.add_argument("--end", type=int, default=None)
	parser.add_argument("--mid", type=int, default=None)
	parser.add_argument("--delta", type=int, default=100)
	parser.add_argument("--strategy", choices=["brutal", "jacobi", "both"], default="jacobi")
	args = parser.parse_args()
	plot(args.start, args.end, args.mid, args.delta, args.strategy)