import sympy as sym
import time
import matplotlib.pyplot as plt

def descartes(ax, ran_x, ran_y, ax_title,x_label = "x", y_label = "y"):
    ax.set_xlabel(x_label, fontsize = 12)
    ax.set_ylabel(y_label, fontsize = 12)
    ax.set_xlim(ran_x[0], ran_x[1])
    ax.set_ylim(ran_y[0], ran_y[1])
    ax.set_title(ax_title, fontsize = 14)
    ax.grid()
    ax.axhline(0, color = "black")
    ax.axvline(0, color = "black")

fig = plt.figure(figsize = (10, 10))
ax = fig.add_subplot(111)
ab=10
nn=10

title1 = "Tokyo University 2024 Math qu.6 NO2"
descartes(ax, [-1*ab,ab], [-1*ab, ab],title1)

start=time.time()

i=0
for a in range(-1*ab,ab):
    for b in range(-1*ab,ab):
        for n in range(-1*nn,nn):
            fn=n**3+a*n**2+b*n
            if sym.isprime(fn) == True:
                i += 1
        if i==4:
            ax.scatter(a,b,color="red")
        elif i==3:
            ax.scatter(a,b,color="blue")
        elif i==2:
            ax.scatter(a,b,color="green")
        elif i==1:
            ax.scatter(a,b,color="yellow")
        else:
            ax.scatter(a,b,color="pink")
        i=0

end=time.time()
run_time=end-start
print("実行時間=",run_time)
plt.show()