'''
observation:
 ([カートの位置，カートの速度，ポールの角度，ポールの角速度]，reward, 終了条件×２，{Info})
(array([ 0.017, -0.01 ,  0.002, -0.01 ], dtype=float32), {})

rewardは常に1．rewardの累計を大きくするのが目標

'''
import gymnasium as gym
import numpy as np
import time

env=gym.make('CartPole-v1',render_mode='human')
total_reward = 0
(state,_)=env.reset()


def now_policy(state):
    time.sleep(0.1)  # Sleep for 100ms to slow down the rendering
    state = int(np.random.randint(0, 2))
    if state >= 1:
        action = 1
    else:
        action = 0
    print(state)
    return action

episode_over = False
while not episode_over:
    action = now_policy(state)
    observation, reward, terminated, truncated, info = env.step(action)
    episode_over = terminated or truncated
    print(observation, reward, terminated, truncated, info)
    total_reward += reward

print(total_reward)
env.close()

# Q値を出力するためのニューラルネットワーク
class N_Network(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(parameters["input_size"], 128) # 入力4次元 → 隠れ層128ユニットへの全結合層（例：観測が4次元のとき）
        self.l2 = nn.Linear(128, 128)
        self.l3 = nn.Linear(128, parameters["action_size"])

    # この関数を実行すれば深層学習が出来る
    def forward(self, x):
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = self.l3(x)
        return x

class ReplayBuffer:
    # バッファーの定義，バッチサイズの代入
    def __init__(self, buffer_size, batch_size):
        self.batch_size = batch_size
        self.buffer_size = buffer_size
        self.buffer = deque(maxlen=buffer_size) # dequeはリストのようなもの

    #関数によっては__len__()という関数が内部で知らないうちに使われる)
    def __len__(self):
        return len(self.buffer)

    # フィードバック[カートの位置，カートの速度，ポールの角度，ポールの角速度]を取得しリプレイバッファに保存．
    # 溢れたら古いデータを削除．ここでの経験は[状態，行動，報酬，次の状態]の4つ
    def add(self, experience):
        self.buffer.append(experience)

    # バッファーからランダムにバッチサイズ分のデータを取得
    def sampling(self):
        datas = np.random.choice(self.buffer, self.batch_size, replace=False) # batch_size分のデータをランダムに取得
        
        # データを分割してそれぞれの変数に格納
        state = torch.tensor(np.stack([x[0] for x in datas]))
        action = torch.tensor(np.array([x[1] for x in datas]).astype(np.int64))
        reward = torch.tensor(np.array([x[2] for x in datas]).astype(np.float32))
        next_state = torch.tensor(np.stack([x[3] for x in datas]))
        done = torch.tensor(np.array([x[4] for x in datas]).astype(np.int32))
        return state, action, reward, next_state, done

# ε-greedyポリシーに従って行動を選択
def select_action(self, state):
    if np.random.rand() <= self.epsilon:
        # εの確率でランダムな行動を選択
        return self.env.action_space.sample()
    
    # 1-εの確率で最適方策に基づいて価値関数を更新
    else:
        state = torch.FloatTensor(state).unsqueeze(0) # PyTorchのTensorに変換＋バッチ形式に対応
        with torch.no_grad(): # 勾配計算を無効化
            return self.q_net(state).argmax().item()

def calc_TDerror(states, actions, next_states, rewards, dones):
    # QネットワークとターゲットネットワークからQ値を取得（TD誤差を求めるため）
    actions = torch.LongTensor(actions).unsqueeze(1)  # [batch_size] → [batch_size, 1]
    q_values = Qnet(states) # → Q(s_t, ·) 全アクションのQ値を出力（形状：[batch_size, action_dim]）
    # print(q_values, action)
    q_value = q_values.gather(1, actions).squeeze() # actionに対応するQ値を出力（形状：[batch_size]

    with torch.no_grad():
        next_q_max = Tnet(next_states).max(1)[0]
        q_value_target = rewards + (1 - dones) * parameters["gamma"] * next_q_max # Q値のターゲット値をベルマン方程式 に従って計算

    # Q値から価値関数とTD誤差を計算．TD誤差を最小化するようにQネットワークのパラメータを更新する
    loss = torch.nn.MSELoss()(q_value, q_value_target)

    # 深層学習のルールに従ってQネットワークのパラメータを更新
    optimizer.zero_grad() # 勾配を初期化
    loss.backward() # 誤差逆伝播法で勾配を計算
    optimizer.step() # Qネットワークのパラメータ(重み)を更新