from keras import backend as K
from keras.layers import Dense, Input, Flatten, Conv2D
from keras.models import Model, load_model
from keras.optimizers import Adam
from keras.utils import plot_model

verbose = 0
action_dim = env.action_space.n
action_space = [i for i in range(action_dim)] # ['NOOP', 'FIRE', 'RIGHT', 'LEFT']
discount = 0.97
actor_lr = 0.001 #actorの学習率
critic_lr = 0.001 #criticの学習率
pretrain_use = 0 #1にすると訓練されたモデルを使う
actor_h5_path = 'atari_ac_actor.h5'
critic_h5_path = 'atari_ac_critic.h5'

#モデル構築
input = Input(shape=(frame_length, frame_width, frame_height))
delta = Input(shape=[1])
con1 = Conv2D(32, (8, 8), strides=(4, 4), padding='same', activation='relu')(input)
con2 = Conv2D(64, (4, 4), strides=(2, 2), padding='same', activation='relu')(con1)
fla1 = Flatten()(con2)
dense = Dense(128, activation='relu')(fla1) #prob, valueをシェア
prob = Dense(action_dim, activation='softmax')(dense) #actor部分
value = Dense(1, activation='linear')(dense) #critic部分

#ロス関数の定義
def custom_loss(y_true, y_pred):
    out = K.clip(y_pred, 1e-8, 1-1e-8) #限界を設定
    log_lik = y_true * K.log(out) #方策勾配
    return K.sum(-log_lik * delta)

if pretrain_use == 1:
    #訓練されたモデルを使う
    actor = load_model(actor_h5_path, custom_objects={'custom_loss': custom_loss}, compile=False)
    critic = load_model(critic_h5_path)

actor = Model(inputs=[input, delta], outputs=[prob])
critic = Model(inputs=[input], outputs=[value])
policy = Model(inputs=[input], outputs=[prob])

actor.compile(optimizer=Adam(lr=actor_lr), loss=custom_loss)
critic.compile(optimizer=Adam(lr=critic_lr), loss='mean_squared_error')

#アクションを予測
def action_choose(state, epsilon, episode, action_space):
    #epsilonは最初に1に設定して徐々に下げる
    #毎回行動する時ランダム数字と比べて
    #epsilonの方が大きいならランダムアクションを取る
    if epsilon >= random.random() or episode < initial_replay:
        action = random.randrange(action_dim)
    else:
        probabiliy = policy.predict(state)[0]
        #予測した結果は4つのアクションに対してそれぞれ確率がある
        #その確率に沿ってアクションを選ぶ
        action = np.random.choice(action_space, p=probabiliy)
    return action

#データを学習
def network_learn(state, action, reward, state_next, done):
    reward_clip = np.sign(reward)
    critic_value = critic.predict(state)
    critic_value_next = critic.predict(state_next)

    target = reward_clip + discount * critic_value_next * (1 - int(done))
    delta =  target - critic_value
    actions = np.zeros([1, action_dim])
    actions[np.arange(1), action] = 1

    actor.fit([state, delta], actions, verbose=verbose)
    critic.fit(state, target, verbose=verbose)