import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
from matplotlib.colors import ListedColormap

#枠なしの座標を準備
fig = plt.figure()
ax = fig.add_subplot(111)
ax.axis("off")
ax.set_aspect('equal')
ax.set_xlim(  0, 50)
ax.set_ylim(-20,  0)
ax.set_title("2021一橋大数学　前期第１問",pad=50)

#初期配列の作成
numbers = np.array(range(1,1001))
isPrime = np.ones(1000,dtype=bool)

x = np.arange( 0, 51,  1)
y = np.arange( 0,-21, -1)
X,Y = np.meshgrid(x,y)

imgs=[]

#１は素数じゃない
isPrime[0] = False
meshPrime = np.reshape(isPrime,(20,50)).astype(int)
meshPrime[0][0] = 2

ta1 = ax.text(20,2,str(1),fontsize=20,ha="right",va="bottom")
ta2 = ax.text(21,2,"は素数ではありません！",fontsize=14,ha="left",va="bottom")
tb1 = ax.text(42.5,-24,"残り",ha="right",va="bottom")
tb2 = ax.text(45.0,-24,str(1000),fontsize=18,ha="center",va="bottom")
tb3 = ax.text(47.5,-24,"個",ha="left",va="bottom")

for i in range(15):
    cmap0 = ListedColormap(["grey",(1,0.5,0,1-i/14)])
    mesh0  = ax.pcolormesh(X,Y,meshPrime,cmap=cmap0)
    img = [mesh0]+[ta1]+[ta2]+[tb1]+[tb2]+[tb3]
    imgs.append(img)

#エラトステネスのふるい
prime = 2
while prime<=997:
    
    meshPrime = np.reshape(isPrime,(20,50)).astype(int)
    meshPrime[(prime-1)//50][(prime-1)%50] = 2
    
    ta1 = ax.text(20,2,str(prime),fontsize=20,ha="right",va="bottom")
    ta2 = ax.text(21,2,"は素数",fontsize=14,ha="left",va="bottom")
    tb2 = ax.text(45.0,-24,'{:>3}'.format(len(numbers[isPrime])),fontsize=18,ha="center",va="bottom")
    
    for i in range(5):
        cmap1 = ListedColormap(["white","grey",(1,0.5,0)])
        mesh1 = ax.pcolormesh(X,Y,meshPrime,cmap=cmap1)
        img = [mesh1]+[ta1]+[ta2]+[tb1]+[tb2]+[tb3]
        imgs.append(img)
    
    toDelete = numbers%prime == 0
    toDelete[prime-1] = False
    
    if prime<=31:
        meshPrime  = 2*np.reshape(isPrime,(20,50)).astype(int)
        meshDelete = np.reshape(toDelete,(20,50)).astype(int)
        meshtoShow = meshPrime - meshDelete
        meshtoShow[(prime-1)//50][(prime-1)%50] = 3
            
        ta2 = ax.text(21,2,"の倍数は脱落！",fontsize=14,ha="left",va="bottom")
        
        for i in range(15):
            if prime == 2:
                cmap2 = ListedColormap(["white",(1,0,0,1-i/14),"grey",(1,0.5,0)])
            else:
                cmap2 = ListedColormap([(1,0,0,0.2*(1-i/14)),"white",(1,0,0,1-i/14),"grey",(1,0.5,0)])
            mesh2 = ax.pcolormesh(X,Y,meshtoShow,cmap=cmap2)
            img = [mesh2]+[ta1]+[ta2]+[tb1]+[tb2]+[tb3]
            imgs.append(img)
    
    isPrime = isPrime & ~toDelete
    prime = prime+np.argmin(~isPrime[prime:])+1
    
mov=ani.ArtistAnimation(fig, imgs, 200)
plt.show()