"""Example: train an agent with Stable-Baselines3 using action masking.

Requires: pip install blokus-gym[train]
"""
from stable_baselines3 import PPO

from blokus_gym import ActionMaskWrapper, BlokusEnv


def main():
    # Create environment with ActionMaskWrapper
    env = BlokusEnv(num_players=2, board_size=7, bot_type=None)
    env = ActionMaskWrapper(env)

    # Train PPO agent
    model = PPO(
        "MultiInputPolicy",
        env,
        verbose=1,
        learning_rate=3e-4,
        n_steps=2048,
        batch_size=64,
        n_epochs=10,
        gamma=0.99,
        seed=42,
    )

    print("Training for 10,000 steps...")
    model.learn(total_timesteps=10000)

    # Save the model
    model.save("blokus_ppo_agent")

    # Evaluate
    obs, info = env.reset(seed=42)
    total_reward = 0
    steps = 0
    terminated = False
    truncated = False

    while not (terminated or truncated):
        action, _ = model.predict(obs, deterministic=True)
        obs, reward, terminated, truncated, info = env.step(action)
        total_reward += reward
        steps += 1

    print(f"Evaluation: {steps} steps, reward: {total_reward}")

    # Cleanup
    env.close()


if __name__ == "__main__":
    main()