"""Unified training entrypoint for PeptideGym."""
from __future__ import annotations

import json
import os
from typing import Any


def train_all(
    envs: list[str],
    steps_per_env: int = 90_000,
    seed: int = 42,
    output_path: str = "results/training_results.json",
) -> list[dict[str, Any]]:
    """Train on multiple environments and save aggregated results."""
    from peptidegym.agents.ppo import train_ppo

    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)

    all_results: list[dict[str, Any]] = []
    for env_id in envs:
        print(f"Training on {env_id} for {steps_per_env} steps ...")
        result = train_ppo(env_id, total_timesteps=steps_per_env, seed=seed)
        all_results.append(result)

    with open(output_path, "w") as f:
        json.dump(all_results, f, indent=2)

    print(f"Results saved to {output_path}")
    return all_results


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description="Train PeptideGym agents")
    parser.add_argument(
        "--envs",
        nargs="+",
        default=["PeptideGym/AMP-v0"],
        help="Environment IDs to train on",
    )
    parser.add_argument("--steps", type=int, default=90_000)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", default="results/training_results.json")
    args = parser.parse_args()

    results = train_all(args.envs, args.steps, args.seed, args.output)
    for r in results:
        print(json.dumps(r, indent=2))


if __name__ == "__main__":
    main()
