#!/usr/bin/env python3
"""train.py — PPO over a population of bodies (the embodiment-scaling sweep).

  python3 train.py --mode scaling --n_bodies 1    --steps 100_000_000 --tag n1
  python3 train.py --mode scaling --n_bodies 10   --steps 100_000_000 --tag n10
  python3 train.py --mode scaling --n_bodies 100  --steps 100_000_000 --tag n100
  python3 train.py --mode scaling --n_bodies 1000 --steps 100_000_000 --tag n1000
  python3 train.py --mode single  --steps 50_000_000 --tag specialist_hex     # one body
  python3 train.py --mode test:quad_long --steps 50_000_000 --tag specialist_quad_long

Smoke test on CPU (minutes, learns nothing): python3 train.py --smoke
GPU: num_envs 4096 is the normal setting; a ~1k-body sweep run is a few hours on one GPU.
"""
import argparse, functools, json, os, time, pickle
import compat  # noqa: F401  (jax/brax API shim)
import jax, jax.numpy as jp, numpy as np
from brax.training.agents.ppo import train as ppo
from brax.training.agents.ppo import networks as ppo_networks
import env as E

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", default="scaling", help="scaling | single | test:<name>")
    ap.add_argument("--n_bodies", type=int, default=100)
    ap.add_argument("--steps", type=int, default=100_000_000)
    ap.add_argument("--num_envs", type=int, default=4096)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--tag", default="run")
    ap.add_argument("--smoke", action="store_true")
    a = ap.parse_args()
    if a.smoke:   # tiny: proves the loop runs, learns nothing
        a.steps, a.num_envs, a.n_bodies = 4 * 4 * 8 * 2, 4, 4

    env = E.ModBotEnv()
    rand = E.make_randomizer(env, mode=a.mode, n_bodies=a.n_bodies, seed=a.seed)
    os.makedirs("runs", exist_ok=True)
    log = []

    def progress(step, metrics):
        row = {"step": int(step), "reward": float(metrics.get("eval/episode_reward", np.nan)),
               "ep_len": float(metrics.get("eval/avg_episode_length", np.nan)),
               "r_lin": float(metrics.get("eval/episode_r_lin", np.nan)),
               "fall": float(metrics.get("eval/episode_fall", np.nan)), "t": time.time() - t0}
        log.append(row)
        print(f"  {row['step']:>12,d}  reward {row['reward']:8.2f}  len {row['ep_len']:6.0f}  "
              f"lin {row['r_lin']:6.1f}  falls {row['fall']:5.2f}  {row['t']:6.0f}s", flush=True)

    hid = (256, 256, 128) if not a.smoke else (32, 32)
    kw = dict(num_timesteps=a.steps, num_evals=10 if not a.smoke else 1, reward_scaling=1.0,
              episode_length=env.max_steps if not a.smoke else 50, normalize_observations=True,
              action_repeat=1, unroll_length=20 if not a.smoke else 4,
              num_minibatches=32 if not a.smoke else 1, num_updates_per_batch=4 if not a.smoke else 1,
              discounting=0.97, learning_rate=3e-4, entropy_cost=1e-2, num_envs=a.num_envs,
              batch_size=256 if not a.smoke else 4, seed=a.seed, randomization_fn=rand,
              num_eval_envs=128 if not a.smoke else 4,
              network_factory=functools.partial(ppo_networks.make_ppo_networks,
                                                policy_hidden_layer_sizes=hid,
                                                value_hidden_layer_sizes=hid))
    print(f"mode={a.mode} bodies={a.n_bodies} envs={a.num_envs} steps={a.steps:,}  on {jax.devices()[0].platform}")
    t0 = time.time()
    make_inference_fn, params, metrics = ppo.train(environment=env, progress_fn=progress, **kw)
    out = f"runs/{a.tag}"
    with open(out + ".pkl", "wb") as f:
        pickle.dump({"params": params, "mode": a.mode, "n_bodies": a.n_bodies, "seed": a.seed, "hidden": hid}, f)
    json.dump(log, open(out + ".json", "w"), indent=1)
    print(f"saved {out}.pkl  ({time.time()-t0:.0f}s)")

if __name__ == "__main__":
    main()