import numpy as np
import matplotlib.pyplot as plt

DNA_SIZE = 1             #   DNAの数
DNA_BOUND = [-5, 7]      #   解を探す範囲の指定
N_GENERATIONS = 200
MUT_STRENGTH = 5.        #　　　初期ステップサイズ　動的突然変異強度

#最大値を求める関数の定義
def F(x): return -(x*x*x*x) + 3*(x*x*x) + 17*(x*x) - 8*x


# find non-zero fitness for selection
def get_fitness(pred): return pred.flatten()


def make_kid(parent):    #子の計算
    # 交叉はなしで突然変異のみとする
    k = parent + MUT_STRENGTH * np.random.randn(DNA_SIZE)
    k = np.clip(k, *DNA_BOUND)
    return k


def kill_bad(parent, kid):   #親と子の比較
    global MUT_STRENGTH
    fp = get_fitness(F(parent))[0]
    fk = get_fitness(F(kid))[0]
    p_target = 1/5
    if fp < fk:     # 子が親よりベター
        parent = kid
        ps = 1.     # kid win -> ps = 1 (成功した子)
    else:
        ps = 0.
    # グローバル突然変異強度の調整
    MUT_STRENGTH *= np.exp(1/np.sqrt(DNA_SIZE+1) * (ps - p_target)/(1 - p_target))
    return parent


parent = 5 * np.random.rand(DNA_SIZE)   # 親のDNA

plt.ion()       # 作図
x = np.linspace(*DNA_BOUND, 200)


for _ in range(N_GENERATIONS):
    # ES part
    kid = make_kid(parent)
    py, ky = F(parent), F(kid)
    parent = kill_bad(parent, kid)   #

    # 作図
    plt.cla()
    plt.xlabel('x')
    plt.ylabel('f(x)')
    plt.scatter(parent, py, s=200, lw=0, c='red', alpha=0.5,)
    plt.scatter(kid, ky, s=200, lw=0, c='blue', alpha=0.5)
    plt.text(-2, -500, 'Mutation strength=%.2f' % MUT_STRENGTH)
    plt.plot(x, F(x)); plt.pause(0.05)
plt.ioff(); plt.show()
print ('at x=',parent, 'f(x)max=',py)