""" reinforcement learning，在监督学习之后，进行自我对弈的强化学习的策略网络 """
import torch
from tqdm import tqdm

from tugo.agents.policy_gradient_agent import PolicyAgent
from tugo.agents.predict import DeepLearningAgent
from tugo.encoders.alphago import AlphaGoEncoder
from tugo.rl.simulate import experience_simulation

encoder = AlphaGoEncoder()

# Load opponents
sl_agent = DeepLearningAgent.load('checkpoints/alphago_sl_policy.pt', encoder)
sl_opponent = DeepLearningAgent.load('checkpoints/alphago_sl_policy.pt', encoder)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

alphago_rl_agent = PolicyAgent(sl_agent.model.to(device), encoder)
opponent = PolicyAgent(sl_opponent.model.to(device), encoder)

# Run simulation
num_games = 1000

with tqdm(total=num_games) as pbar:
    experience = experience_simulation(num_games, alphago_rl_agent, opponent, progress_callback=lambda: pbar.update(1))

alphago_rl_agent.train(experience)

# # Serialize RL agent
# alphago_rl_agent.save('checkpoints/alphago_rl_policy.pt')

# # Serialize experience
# experience.save('checkpoints/alphago_rl_experience.pt')



def generate_incremental_filepath(file_path):
    import os
    directory, filename = os.path.split(file_path)
    filename_base, file_extension = os.path.splitext(filename)

    # 查找保存目录下的最后一个文件序号
    last_index = 0
    file_exists = False

    for file_name in os.listdir(directory):
        if file_name.startswith(filename_base) and file_name.endswith(file_extension):
            try:
                index = int(file_name[len(filename_base):-len(file_extension)])
                last_index = max(last_index, index)
                file_exists = True
            except ValueError:
                continue

    if not file_exists:
        last_index = 0

    # 生成带序号的新文件名
    new_index = last_index + 1
    new_filename = f"{filename_base}_{new_index}{file_extension}"

    # 生成新文件的完整路径
    new_file_path = os.path.join(directory, new_filename)

    return new_file_path


# Serialize RL agent
alphago_rl_agent.save(generate_incremental_filepath('checkpoints/alphago_rl_policy.pt'))

# Serialize experience
experience.save(generate_incremental_filepath('checkpoints/alphago_rl_experience.pt'))