import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches
import numpy as np

Source = [(0, 1), (0, 0)]
Right = [(0, 0), (1, 0), (2, 0), (2, -1), (2, -2), (2, -3), (1, -3), (0, -3)]
Left = [(0, 0), (-1, 0), (-2, 0), (-3, 0), (-4, 0), (-5, 0), (-5, -1),
        (-5, -2), (-5, -3), (-4, -3), (-3, -3), (-2, -3), (-1, -3), (0, -3)]
Target = [(0, -3), (0, -4)]

Route = [Source, Right, Left, Target]
tau = np.ones([10, 20])*0.1
eta = np.ones([10, 20])
plt.ion()
fig, ax = plt.subplots()
plt.pause(3)


def plt_reset(init=0):
    plt.cla()
    plt.style.use('fivethirtyeight')
    plt.axis('off')  # 去掉坐标轴
    plt.axis('equal')
    ax.set_xlim(-10, 10)
    ax.set_ylim(-10, 10)
    ax.annotate('Source', xy=(0.5, 1.5), xytext=(3, 4),
                arrowprops=dict(facecolor='black', shrink=0.05))
    if init:
        plt.pause(init)
    ax.annotate('Target', xy=(0.5, -3.5), xytext=(3, -6),
                arrowprops=dict(facecolor='black', shrink=0.05))
    if init:
        plt.pause(init)


plt_reset(1)

for v in Source:
    patch = patches.Rectangle(v, 0.9, 0.9, color=(0.9, 0.9, 1))
    ax.add_patch(patch)
    plt.pause(0.1)
for v in Target:
    patch = patches.Rectangle(v, 0.9, 0.9, color=(0.9, 0.9, 1))
    ax.add_patch(patch)
    plt.pause(0.1)
for v in Right:
    patch = patches.Rectangle(v, 0.9, 0.9, color=(0.9, 0.9, 1))
    ax.add_patch(patch)
    plt.pause(0.1)
for v in Left:
    patch = patches.Rectangle(v, 0.9, 0.9, color=(0.9, 0.9, 1))
    ax.add_patch(patch)
    plt.pause(0.1)

start_idx = [[1, 0], [2, 0]]
ant_idx = [[1, 0], [2, 0]]
plt.pause(1)
for _ in range(1000):
    plt_reset()
    for i, r in enumerate(Route):
        for j, v in enumerate(r):
            tau[i][j] *= 0.9
            t = tau[i][j]
            patch = patches.Rectangle(
                v, 0.9, 0.9, color=(np.clip(1-t, 0, 1), np.clip(1-t, 0, 1), 1))
            ax.add_patch(patch)

    # plt.pause(0.1)

    for i in range(len(ant_idx)):
        # 更新蚂蚁坐标
        x, y = ant_idx[i]
        patch = patches.Rectangle(Route[x][y], 0.5, 0.5, color=(0, 0, 0))
        ax.add_patch(patch)

        tau[x][y] += 1

    plt.pause(0.001)

    def Pk_ij(allowed, alpha=1, beta=0.1):
            # 回到起点
        top = np.power(tau[allowed], alpha)*np.power(eta[allowed], beta)
        bottom = np.sum(np.power(tau[allowed], alpha)
                        * np.power(eta[allowed], beta))
        return top/bottom

    for i in range(len(ant_idx)):
        x, y = ant_idx[i]
        if y == len(Route[x]) - 1:
            if x == 3:
                Pk = Pk_ij(((1, 2), (1, 1)))
                j = np.argmax(Pk)
                ant_idx[i] = start_idx[j].copy()
            else:
                ant_idx[i] = [3, 0]
        else:
            ant_idx[i][1] += 1


plt.pause(10)