"""PPO training wrapper for PeptideGym environments.

Requires the ``train`` optional dependency group: ``pip install peptidegym[train]``
"""
from __future__ import annotations

import json
import os
from typing import Any

import numpy as np


def train_ppo(
    env_id: str,
    total_timesteps: int = 90_000,
    seed: int = 42,
    log_dir: str = "results/training_logs",
) -> dict[str, Any]:
    """Train PPO on a PeptideGym environment. Returns metrics dict."""
    # Import SB3 here to keep it optional
    from stable_baselines3 import PPO
    from stable_baselines3.common.vec_env import DummyVecEnv
    import gymnasium as gym
    import peptidegym  # noqa: F401 — triggers env registration
    from peptidegym.envs.wrappers import FlattenObservation

    os.makedirs(log_dir, exist_ok=True)

    def _make_env() -> gym.Env:
        return FlattenObservation(gym.make(env_id))

    env = DummyVecEnv([_make_env])
    model = PPO(
        "MlpPolicy",
        env,
        verbose=1,
        seed=seed,
        learning_rate=3e-4,
        n_steps=2048,
        batch_size=64,
        tensorboard_log=log_dir,
    )
    model.learn(total_timesteps=total_timesteps)

    # Evaluate
    eval_env = FlattenObservation(gym.make(env_id))
    rewards: list[float] = []
    for _ in range(100):
        obs, _ = eval_env.reset()
        total_reward = 0.0
        done = False
        while not done:
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, _ = eval_env.step(action)
            total_reward += float(reward)
            done = terminated or truncated
        rewards.append(total_reward)
    eval_env.close()

    return {
        "mean_reward": float(np.mean(rewards)),
        "std_reward": float(np.std(rewards)),
        "env_id": env_id,
        "total_timesteps": total_timesteps,
    }


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description="Train PPO on a PeptideGym env")
    parser.add_argument("--env", default="PeptideGym/AMP-v0")
    parser.add_argument("--steps", type=int, default=90_000)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    results = train_ppo(args.env, args.steps, args.seed)
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()
