import matplotlib
import numpy as np
from numpy.ma import cos
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import datetime

matplotlib.use('Qt5Agg')

DNA_SIZE = 24  # 编码长度
POP_SIZE = 100  # 种群大小
CROSS_RATE = 0.8  # 交叉率
MUTA_RATE = 0.15  # 变异率
Iterations = 1000  # 代次数
X_BOUND = [0, 10]  # X区间
Y_BOUND = [0, 10]  # Y区间


def F(x, y):  # 适应度函数
    return (6.452 * (x + 0.125 * y) * (cos(x) - cos(2 * y)) ** 2) / (
            0.8 + (x - 4.2) ** 2 + 2 * (y - 7) ** 2) + 3.226 * y


def decodeDNA(pop):  # 解码
    x_pop = pop[:, 1::2]  # 奇数列表示X
    y_pop = pop[:, ::2]  # 偶数列表示y
    x = x_pop.dot(2 ** np.arange(DNA_SIZE)[::-1]) / float(2 ** DNA_SIZE - 1) * (X_BOUND[1] - X_BOUND[0]) + X_BOUND[0]
    y = y_pop.dot(2 ** np.arange(DNA_SIZE)[::-1]) / float(2 ** DNA_SIZE - 1) * (Y_BOUND[1] - Y_BOUND[0]) + Y_BOUND[0]
    return x, y


def getfitness(pop):
    x, y = decodeDNA(pop)
    temp = F(x, y)
    return (temp - np.min(temp)) + 0.0001  # 减去最小的适应度是为了防止适应度出现负数


def select(pop, fitness):  # 根据适应度选择
    temp = np.random.choice(np.arange(POP_SIZE), size=POP_SIZE, replace=True, p=(fitness) / (fitness.sum()))
    return pop[temp]


def crossmuta(pop, CROSS_RATE):
    new_pop = []
    for i in pop:  # 遍历种群中的每一个个体，将该个体作为父代
        temp = i  # 子代先得到父亲的全部基因
        if np.random.rand() < CROSS_RATE:  # 以交叉概率发生交叉
            j = pop[np.random.randint(POP_SIZE)]  # 从种群中随机选择另一个个体，并将该个体作为母代
            cpoints1 = np.random.randint(0, DNA_SIZE * 2 - 1)  # 随机产生交叉的点
            cpoints2 = np.random.randint(cpoints1, DNA_SIZE * 2)
            temp[cpoints1:cpoints2] = j[cpoints1:cpoints2]  # 子代得到位于交叉点后的母代的基因
        mutation(temp, MUTA_RATE)  # 后代以变异率发生变异
        new_pop.append(temp)
    return new_pop


def mutation(temp, MUTA_RATE):
    if np.random.rand() < MUTA_RATE:  # 以MUTA_RATE的概率进行变异
        mutate_point = np.random.randint(0, DNA_SIZE)  # 随机产生一个实数，代表要变异基因的位置
        temp[mutate_point] = temp[mutate_point] ^ 1  # 将变异点的二进制为反转


def print_info(pop):  # 用于输出结果
    fitness = getfitness(pop)
    maxfitness = np.argmax(fitness)  # 返回最大值的索引值
    print("max_fitness:", fitness[maxfitness])
    x, y = decodeDNA(pop)
    print("最优的基因型：", pop[maxfitness])
    print("(x, y):", (x[maxfitness], y[maxfitness]))
    print("F(x,y)_max = ", F(x[maxfitness], y[maxfitness]))


def plot_3d(ax):
    X = np.linspace(*X_BOUND, 100)
    Y = np.linspace(*Y_BOUND, 100)
    X, Y = np.meshgrid(X, Y)
    Z = F(X, Y)
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm)
    ax.set_zlim(-20, 100)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    plt.pause(0.01)  # 缩短暂停时间
    # 移除 plt.show()


start_t = datetime.datetime.now()
if __name__ == "__main__":
    plt.ion()  # 提前设置交互模式
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')  # 修改3D坐标创建方式
    plot_3d(ax)

    pop = np.random.randint(2, size=(POP_SIZE, DNA_SIZE * 2))
    for _ in range(Iterations):
        x, y = decodeDNA(pop)
        if 'sca' in locals():
            sca.remove()
        sca = ax.scatter(x, y, F(x, y), c='red', marker='o', s=50)  # 增大标记尺寸
        plt.draw()  # 强制重绘
        plt.pause(0.01)  # 保证足够的更新时间
        pop = np.array(crossmuta(pop, CROSS_RATE))
        fitness = getfitness(pop)
        pop = select(pop, fitness)  # 选择生成新的种群
end_t = datetime.datetime.now()
print((end_t - start_t).seconds)
print_info(pop)
plt.ioff()
plot_3d(ax)