#!/usr/bin/env python3
"""
UGTC-PPO on Procgen Benchmark (hard mode) — 25M environment steps.

Default configuration matches the paper:
    - env:               coinrun (hard mode, 500 training levels)
    - evaluation:        500 unseen levels from level 10,000+
    - total timesteps:   25,000,000
    - vectorized envs:   64

UGTC hyperparameters (fixed across ALL benchmarks):
    λ_fast = 0.80   (fast critic GAE lambda)
    λ_slow = 0.99   (slow ensemble GAE lambda)
    M = 3           (ensemble size)
    β = 5.0         (gate temperature)
    α_EMA = 0.99    (EMA momentum)

This script is derived from the benchmark code used in the paper experiments.
Results may vary by hardware, random seed, and library version.

Requirements:
    pip install torch procgen gymnasium pandas psutil

Usage:
    python benchmarks/procgen/train_ugtc_ppo_procgen.py
    python benchmarks/procgen/train_ugtc_ppo_procgen.py --env_name bigfish
"""

import json
import math
import os
import random
import time
import argparse
from pathlib import Path
from typing import Any, Dict

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.categorical import Categorical

import gymnasium as gym

try:
    import procgen2  # noqa
except ImportError:
    try:
        import procgen  # noqa
    except ImportError:
        raise ImportError("Install procgen or procgen2: pip install procgen")

try:
    import pandas as pd
    HAS_PANDAS = True
except ImportError:
    HAS_PANDAS = False

try:
    import psutil
    HAS_PSUTIL = True
except ImportError:
    HAS_PSUTIL = False


# ── Default config ────────────────────────────────────────────────────────────

DEFAULT_CONFIG: Dict[str, Any] = {
    "env_name": "coinrun",
    "distribution_mode": "hard",
    "num_levels": 500,
    "start_level": 0,
    "eval_start_level": 10000,
    "eval_num_levels": 500,
    "seed": 1,
    "total_timesteps": 25_000_000,
    "num_envs": 64,
    "num_steps": 256,
    "update_epochs": 3,
    "num_minibatches": 8,
    "gamma": 0.999,
    # UGTC hyperparameters (fixed across all benchmarks)
    "gae_lambda_fast": 0.80,
    "gae_lambda_slow": 0.99,
    "slow_critics": 3,
    "uncertainty_beta": 5.0,
    "running_unc_momentum": 0.99,
    # PPO hyperparameters
    "clip_coef": 0.2,
    "ent_coef": 0.01,
    "vf_coef_fast": 0.5,
    "vf_coef_slow": 0.5,
    "max_grad_norm": 0.5,
    "learning_rate": 5e-4,
    "anneal_lr": True,
    "feature_dim": 256,
    "hidden_size": 64,
    # Eval / logging
    "eval_freq": 1_000_000,
    "n_eval_episodes": 20,
    "log_freq": 10,
    "save_model": True,
    "root_dir": "runs_ugtc_ppo_procgen",
    "device": "cuda" if torch.cuda.is_available() else "cpu",
}


# ── Utils ─────────────────────────────────────────────────────────────────────

def set_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def make_procgen_env(env_name, distribution_mode, start_level, num_levels, seed):
    for env_id in [f"procgen:procgen-{env_name}-v0", f"procgen-{env_name}-v0"]:
        try:
            env = gym.make(env_id, start_level=start_level, num_levels=num_levels,
                           distribution_mode=distribution_mode)
            try:
                env.reset(seed=seed)
            except TypeError:
                pass
            return env
        except Exception:
            continue
    raise RuntimeError(f"Cannot create Procgen env '{env_name}'")


def make_vector_envs(cfg, train=True):
    start = cfg["start_level"] if train else cfg["eval_start_level"]
    n_levels = cfg["num_levels"] if train else cfg["eval_num_levels"]
    n_envs = cfg["num_envs"] if train else 1

    def thunk(rank):
        def _init():
            return make_procgen_env(cfg["env_name"], cfg["distribution_mode"],
                                    start_level=start, num_levels=n_levels,
                                    seed=cfg["seed"] + rank)
        return _init

    return gym.vector.SyncVectorEnv([thunk(i) for i in range(n_envs)])


def obs_to_tensor(obs, device):
    x = np.asarray(obs)
    if x.ndim == 3:
        x = x[None, ...]
    return torch.as_tensor(np.transpose(x, (0, 3, 1, 2)).astype(np.float32) / 255.0,
                           dtype=torch.float32, device=device)


# ── Model ─────────────────────────────────────────────────────────────────────

class ResidualBlock(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.block = nn.Sequential(
            nn.ReLU(), nn.Conv2d(ch, ch, 3, 1, 1),
            nn.ReLU(), nn.Conv2d(ch, ch, 3, 1, 1),
        )

    def forward(self, x):
        return x + self.block(x)


class ImpalaBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, 3, 1, 1)
        self.pool = nn.MaxPool2d(3, 2, 1)
        self.res = nn.Sequential(ResidualBlock(out_ch), ResidualBlock(out_ch))

    def forward(self, x):
        return self.res(self.pool(self.conv(x)))


class VisualEncoder(nn.Module):
    def __init__(self, feature_dim):
        super().__init__()
        self.net = nn.Sequential(
            ImpalaBlock(3, 16), ImpalaBlock(16, 32), ImpalaBlock(32, 32),
            nn.ReLU(), nn.Flatten(),
        )
        with torch.no_grad():
            flat = self.net(torch.zeros(1, 3, 64, 64)).shape[1]
        self.proj = nn.Sequential(nn.Linear(flat, feature_dim), nn.ReLU())

    def forward(self, x):
        return self.proj(self.net(x))


class ValueHead(nn.Module):
    def __init__(self, feature_dim, hidden):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(feature_dim, hidden), nn.Tanh(), nn.Linear(hidden, 1),
        )
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.orthogonal_(m.weight, math.sqrt(2))
                nn.init.constant_(m.bias, 0)
        nn.init.orthogonal_(self.net[-1].weight, 1.0)

    def forward(self, feat):
        return self.net(feat).squeeze(-1)


class UGTCPPOAgent(nn.Module):
    def __init__(self, n_actions, feature_dim, hidden, slow_critics):
        super().__init__()
        self.encoder = VisualEncoder(feature_dim)
        self.policy_head = nn.Sequential(
            nn.Linear(feature_dim, hidden), nn.Tanh(), nn.Linear(hidden, n_actions),
        )
        self.fast_value = ValueHead(feature_dim, hidden)
        self.slow_values = nn.ModuleList(
            [ValueHead(feature_dim, hidden) for _ in range(slow_critics)]
        )
        nn.init.orthogonal_(self.policy_head[-1].weight, 0.01)

    def get_features(self, obs):
        return self.encoder(obs)

    def get_policy(self, feat):
        return Categorical(logits=self.policy_head(feat))

    def get_fast_value(self, feat):
        return self.fast_value(feat)

    def get_slow_values(self, feat):
        return torch.stack([h(feat) for h in self.slow_values], dim=0)  # [M, B]

    def get_action_and_values(self, obs, action=None):
        feat = self.get_features(obs)
        dist = self.get_policy(feat)
        if action is None:
            action = dist.sample()
        return action, dist.log_prob(action), dist.entropy(), self.get_fast_value(feat), self.get_slow_values(feat)


# ── GAE ───────────────────────────────────────────────────────────────────────

def compute_gae(rewards, dones, values, next_value, gamma, gae_lambda):
    T, N = rewards.shape
    adv = torch.zeros_like(rewards)
    gae = torch.zeros(N, device=rewards.device)
    for t in reversed(range(T)):
        nxt = next_value if t == T - 1 else values[t + 1]
        nt = 1.0 - (dones[t] if t == T - 1 else dones[t + 1])
        delta = rewards[t] + gamma * nxt * nt - values[t]
        gae = delta + gamma * gae_lambda * nt * gae
        adv[t] = gae
    return adv, adv + values


# ── Eval ──────────────────────────────────────────────────────────────────────

@torch.no_grad()
def evaluate(agent, cfg, device):
    env = make_vector_envs(cfg, train=False)
    returns = []
    for ep in range(cfg["n_eval_episodes"]):
        obs, _ = env.reset(seed=cfg["seed"] + 100000 + ep)
        done = np.array([False])
        ret = 0.0
        while not done[0]:
            obs_t = obs_to_tensor(obs, device)
            feat = agent.get_features(obs_t)
            action = torch.argmax(agent.get_policy(feat).logits, dim=-1)
            out = env.step(action.cpu().numpy())
            obs, reward = out[0], out[1]
            done = np.logical_or(out[2], out[3]) if len(out) == 5 else out[2]
            ret += float(reward[0])
        returns.append(ret)
    env.close()
    return {"eval_mean": float(np.mean(returns)), "eval_std": float(np.std(returns))}


# ── Train ─────────────────────────────────────────────────────────────────────

def train(cfg: Dict[str, Any]) -> None:
    set_seed(cfg["seed"])
    device = torch.device(cfg["device"])
    root = Path(cfg["root_dir"])
    root.mkdir(parents=True, exist_ok=True)

    envs = make_vector_envs(cfg, train=True)
    n_actions = int(envs.single_action_space.n)
    agent = UGTCPPOAgent(n_actions, cfg["feature_dim"], cfg["hidden_size"], cfg["slow_critics"]).to(device)
    optimizer = optim.Adam(agent.parameters(), lr=cfg["learning_rate"], eps=1e-5)

    num_updates = cfg["total_timesteps"] // (cfg["num_envs"] * cfg["num_steps"])
    mb_size = (cfg["num_envs"] * cfg["num_steps"]) // cfg["num_minibatches"]

    # Buffers
    obs_buf    = torch.zeros((cfg["num_steps"], cfg["num_envs"], 3, 64, 64), device=device)
    actions_buf = torch.zeros((cfg["num_steps"], cfg["num_envs"]), dtype=torch.long, device=device)
    logp_buf   = torch.zeros((cfg["num_steps"], cfg["num_envs"]), device=device)
    rew_buf    = torch.zeros((cfg["num_steps"], cfg["num_envs"]), device=device)
    done_buf   = torch.zeros((cfg["num_steps"], cfg["num_envs"]), device=device)
    fv_buf     = torch.zeros((cfg["num_steps"], cfg["num_envs"]), device=device)
    sv_buf     = torch.zeros((cfg["num_steps"], cfg["slow_critics"], cfg["num_envs"]), device=device)

    obs_raw, _ = envs.reset(seed=cfg["seed"])
    next_obs = obs_to_tensor(obs_raw, device)
    next_done = torch.zeros(cfg["num_envs"], device=device)

    global_step = 0
    next_eval = cfg["eval_freq"]
    running_unc = 1.0
    start_time = time.time()
    train_logs, eval_logs = [], []

    print("=" * 80)
    print(f"UGTC-PPO | Procgen {cfg['env_name']} | {cfg['distribution_mode']} | {cfg['total_timesteps']:,} steps")
    print(f"λ_fast={cfg['gae_lambda_fast']}  λ_slow={cfg['gae_lambda_slow']}  M={cfg['slow_critics']}  β={cfg['uncertainty_beta']}")
    print("=" * 80)

    for update in range(1, num_updates + 1):
        if cfg["anneal_lr"]:
            frac = 1.0 - (update - 1) / num_updates
            optimizer.param_groups[0]["lr"] = frac * cfg["learning_rate"]

        for step in range(cfg["num_steps"]):
            global_step += cfg["num_envs"]
            obs_buf[step] = next_obs
            done_buf[step] = next_done
            with torch.no_grad():
                action, logp, _, fv, sv = agent.get_action_and_values(next_obs)
            actions_buf[step], logp_buf[step] = action, logp
            fv_buf[step], sv_buf[step] = fv, sv
            out = envs.step(action.cpu().numpy())
            obs_raw, reward = out[0], out[1]
            terminated, truncated = (out[2], out[3]) if len(out) == 5 else (out[2], out[2])
            rew_buf[step] = torch.as_tensor(reward, dtype=torch.float32, device=device)
            next_obs = obs_to_tensor(obs_raw, device)
            next_done = torch.as_tensor(np.logical_or(terminated, truncated), dtype=torch.float32, device=device)

        with torch.no_grad():
            feat_next = agent.get_features(next_obs)
            nfv = agent.get_fast_value(feat_next)
            nsv = agent.get_slow_values(feat_next).mean(dim=0)
            sv_mean = sv_buf.mean(dim=1)

            adv_fast, ret_fast = compute_gae(rew_buf, done_buf, fv_buf, nfv, cfg["gamma"], cfg["gae_lambda_fast"])
            adv_slow, ret_slow = compute_gae(rew_buf, done_buf, sv_mean, nsv, cfg["gamma"], cfg["gae_lambda_slow"])

            sigma = sv_buf.var(dim=1).sqrt()
            cur_unc = float(sigma.mean().item())
            running_unc = cfg["running_unc_momentum"] * running_unc + (1 - cfg["running_unc_momentum"]) * cur_unc
            gate = torch.sigmoid(-cfg["uncertainty_beta"] * (sigma / (running_unc + 1e-8) - 1.0))
            blended_adv = gate * adv_slow + (1.0 - gate) * adv_fast

        b_obs   = obs_buf.reshape(-1, 3, 64, 64)
        b_act   = actions_buf.reshape(-1)
        b_lp    = logp_buf.reshape(-1)
        b_fret  = ret_fast.reshape(-1)
        b_sret  = ret_slow.reshape(-1)
        b_adv   = blended_adv.reshape(-1)
        b_sv    = sv_buf.permute(0, 2, 1).reshape(-1, cfg["slow_critics"])
        b_gate  = gate.reshape(-1)
        b_adv   = (b_adv - b_adv.mean()) / (b_adv.std() + 1e-8)
        inds    = np.arange(cfg["num_envs"] * cfg["num_steps"])

        for _ in range(cfg["update_epochs"]):
            np.random.shuffle(inds)
            for start in range(0, len(inds), mb_size):
                mb = inds[start:start + mb_size]
                feat = agent.get_features(b_obs[mb])
                dist = agent.get_policy(feat)
                new_lp = dist.log_prob(b_act[mb])
                entropy = dist.entropy()
                fv_new = agent.get_fast_value(feat)
                sv_new = agent.get_slow_values(feat).transpose(0, 1)

                ratio = (new_lp - b_lp[mb]).exp()
                mb_adv = b_adv[mb]
                pg_loss = torch.max(
                    -mb_adv * ratio,
                    -mb_adv * torch.clamp(ratio, 1 - cfg["clip_coef"], 1 + cfg["clip_coef"])
                ).mean()
                fv_loss = 0.5 * ((fv_new - b_fret[mb]) ** 2).mean()
                sv_loss = 0.5 * ((sv_new - b_sret[mb].unsqueeze(-1)) ** 2).mean()
                loss = pg_loss - cfg["ent_coef"] * entropy.mean() + cfg["vf_coef_fast"] * fv_loss + cfg["vf_coef_slow"] * sv_loss
                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                nn.utils.clip_grad_norm_(agent.parameters(), cfg["max_grad_norm"])
                optimizer.step()

        if update % cfg["log_freq"] == 0:
            elapsed = time.time() - start_time
            log = {
                "update": update, "timesteps": global_step,
                "mean_gate": float(b_gate.mean().item()),
                "running_unc": running_unc,
                "sps": int(global_step / max(elapsed, 1e-9)),
            }
            train_logs.append(log)
            print(f"[{update:5d}/{num_updates}] steps={global_step:>9,} gate={log['mean_gate']:.3f} unc={running_unc:.4f} sps={log['sps']}")

        if global_step >= next_eval or update == num_updates:
            ev = evaluate(agent, cfg, device)
            eval_logs.append({"timesteps": global_step, **ev})
            print(f"  ► EVAL  mean={ev['eval_mean']:.3f}  std={ev['eval_std']:.3f}")
            next_eval += cfg["eval_freq"]

    envs.close()

    if HAS_PANDAS:
        pd.DataFrame(eval_logs).to_csv(root / "evals.csv", index=False)
        pd.DataFrame(train_logs).to_csv(root / "train.csv", index=False)

    summary = {
        "algo": "UGTC-PPO", "env_name": cfg["env_name"],
        "lambda_fast": cfg["gae_lambda_fast"], "lambda_slow": cfg["gae_lambda_slow"],
        "M": cfg["slow_critics"], "beta": cfg["uncertainty_beta"],
        "total_timesteps": global_step,
        "best_eval": max((e["eval_mean"] for e in eval_logs), default=float("nan")),
        "elapsed_sec": time.time() - start_time,
    }
    with open(root / "summary.json", "w") as f:
        json.dump(summary, f, indent=2)

    if cfg["save_model"]:
        torch.save({"state_dict": agent.state_dict(), "config": cfg}, root / "model.pt")

    print("\nTraining complete.")
    print(f"  Best eval return: {summary['best_eval']:.3f}")
    print(f"  Results saved to: {root}/")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--env_name", type=str, default="coinrun")
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--total_timesteps", type=int, default=25_000_000)
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()
    cfg = {**DEFAULT_CONFIG, **vars(args)}
    train(cfg)