import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani

fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_xlim(0,5)
ax.set_ylim(0,5)
ax.set_aspect("equal")
ax.axis("off")
ax.grid()
ax.set_title("2020京大数学　理系第５問　文系第５問",pad=15)

#順列でシャッフル
def shuffle(mother):
    
    oldlady = ""
    oldladies = [oldlady]

    for digit in range(len(mother)):
        newladies = []
        
        for oldlady in oldladies:
            for m in mother:
                
                admission = m not in oldlady
                
                if admission:
                    newlady = oldlady + m
                    newladies.append(newlady)
    
        oldladies = newladies
    
    return oldladies

#新しく追加できる行
def get_newrows(k,X):
    
    rows = []
    
    for str_row in shuffle("1234".replace(X[k,0],"")):
        row = np.array(list(str_row))
        
        overlap = False
        for i in range(X[:k,1:].shape[0]):
            if np.any(row==X[i,1:]):
                overlap = True
                
        if overlap:
            pass
        else:
            rows.append(row)
    
    return rows

#一つ一つの例をプロット
def show(M,N,X,count):
    
    img = [ax.text(1.8,4.8,'N = {:>3}'.format(count),
                   ha="left",va="center",color="dimgrey",
                   fontsize=18,fontfamily='fantasy',)]
    
    img += [ax.pcolormesh(N,M,X.astype(int),
                         cmap="Pastel1",edgecolors="w",linewidth=4)]
    
    alphas = [0.2,0.2,0.2,0.2]
    alphas[(count-1)%4] = 1.0
    for i in range(4):
        img += ax.plot(1.9+0.4*i,0.2,alpha=alphas[i],marker="o",color="navy")
        for j in range(4):
            img += [ax.text(i+1,j+1,X[j,i],
                    ha="center",va="center",
                    fontsize=24,fontfamily='fantasy',
                    color="grey")]
            
    if count<=4:
        img += [ax.text(1.6,0.2,"１行目と１列目を固定したとき",ha="right",va="center",color="navy")]
        img += [ax.text(2.5,-0.2,"n="+str(count),ha="center",va="center",
                        color="navy",fontsize=16,fontfamily='fantasy',)]
        img += [ax.axvspan(0.5,4.5,0.7,0.9,
                           fill=False,edgecolor="navy",linewidth=4)]
        img += [ax.axhspan(0.5,3.5,0.1,0.3,
                           fill=False,edgecolor="navy",linewidth=4)]
    elif count<=24:
        img += [ax.text(1.6,0.2,"１行目を固定したとき",ha="right",va="center",color="navy")]
        img += [ax.text(2.5,-0.2,"n="+str(count),ha="center",va="center",
                        color="navy",fontsize=16,fontfamily='fantasy',)]
        img += [ax.axvspan(0.5,4.5,0.7,0.9,
                           fill=False,edgecolor="navy",linewidth=4)]
        
    return img

#1行目と1列目を固定
X = np.zeros((4,4),dtype=str)
Xs = []
for str_row in shuffle("1234"):
    X[0,:] = np.array(list(str_row))
    for str_col in shuffle("1234".replace(X[0,0],"")):
        X[1:,0] = np.array(list(str_col))

        for row in get_newrows(1,X):
            X[1,1:] = row
            for row in get_newrows(2,X):
                X[2,1:] = row
                for row in get_newrows(3,X):
                    X[3,1:] = row
                    Xs.append(X.copy())

M,N = np.mgrid[0.5:5.5:1,0.5:5.5:1]

#結果を表示
count = 0
imgs = []
for i,X in enumerate(Xs):
    count += 1
    X = np.flipud(X)
    img = show(M,N,X,count)
    if i < 4:
        for repeat in range(20):
            imgs.append(img)
    elif i < 24:
        for repeat in range(10):
            imgs.append(img)
    elif count == 576:
        for repeat in range(50):
            imgs.append(img)
    else:
        imgs.append(img)
            
mov=ani.ArtistAnimation(fig, imgs, 100)
plt.show()