"""Unit 8 PPO from scratch (CleanRL-style), updated for Gymnasium / Colab.

Example:
 python ppo.py --env-id LunarLander-v3 --repo-id while0628/ppo-LunarLander-v3-unit8 --total-timesteps 500000 --upload
"""
import argparse
import json
import os
import random
import time
from pathlib import Path

import gymnasium as gym
import imageio
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.categorical import Categorical
from torch.utils.tensorboard import SummaryWriter


def parse_args():
    p = argparse.ArgumentParser(description="PPO for discrete-action Gymnasium environments")
    p.add_argument("--exp-name", default="ppo")
    p.add_argument("--env-id", default="LunarLander-v3")
    p.add_argument("--repo-id", default="while0628/ppo-LunarLander-v3-unit8")
    p.add_argument("--total-timesteps", type=int, default=500_000)
    p.add_argument("--learning-rate", type=float, default=2.5e-4)
    p.add_argument("--num-envs", type=int, default=4)
    p.add_argument("--num-steps", type=int, default=1024)
    p.add_argument("--anneal-lr", action=argparse.BooleanOptionalAction, default=True)
    p.add_argument("--gamma", type=float, default=0.99)
    p.add_argument("--gae-lambda", type=float, default=0.95)
    p.add_argument("--num-minibatches", type=int, default=4)
    p.add_argument("--update-epochs", type=int, default=4)
    p.add_argument("--norm-adv", action=argparse.BooleanOptionalAction, default=True)
    p.add_argument("--clip-coef", type=float, default=0.2)
    p.add_argument("--clip-vloss", action=argparse.BooleanOptionalAction, default=True)
    p.add_argument("--ent-coef", type=float, default=0.01)
    p.add_argument("--vf-coef", type=float, default=0.5)
    p.add_argument("--max-grad-norm", type=float, default=0.5)
    p.add_argument("--target-kl", type=float, default=None)
    p.add_argument("--seed", type=int, default=1)
    p.add_argument("--cuda", action=argparse.BooleanOptionalAction, default=True)
    p.add_argument("--upload", action="store_true", help="Evaluate and upload model/video to Hugging Face Hub")
    p.add_argument("--eval-episodes", type=int, default=10)
    args = p.parse_args()
    args.batch_size = args.num_envs * args.num_steps
    if args.batch_size % args.num_minibatches:
        p.error("num-envs * num-steps must be divisible by num-minibatches")
    args.minibatch_size = args.batch_size // args.num_minibatches
    return args


def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
    nn.init.orthogonal_(layer.weight, std)
    nn.init.constant_(layer.bias, bias_const)
    return layer


class Agent(nn.Module):
    def __init__(self, envs):
        super().__init__()
        obs_dim = int(np.prod(envs.single_observation_space.shape))
        action_dim = envs.single_action_space.n
        self.critic = nn.Sequential(
            layer_init(nn.Linear(obs_dim, 64)), nn.Tanh(),
            layer_init(nn.Linear(64, 64)), nn.Tanh(),
            layer_init(nn.Linear(64, 1), std=1.0),
        )
        self.actor = nn.Sequential(
            layer_init(nn.Linear(obs_dim, 64)), nn.Tanh(),
            layer_init(nn.Linear(64, 64)), nn.Tanh(),
            layer_init(nn.Linear(64, action_dim), std=0.01),
        )

    def get_value(self, x):
        return self.critic(x)

    def get_action_and_value(self, x, action=None):
        logits = self.actor(x)
        probs = Categorical(logits=logits)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action), probs.entropy(), self.critic(x)


def make_env(env_id, seed):
    def thunk():
        env = gym.make(env_id)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        env.action_space.seed(seed)
        return env
    return thunk


def evaluate_and_upload(agent, args, device, model_path, run_name):
    from huggingface_hub import HfApi, get_token

    if get_token() is None:
        raise RuntimeError("No Hugging Face token found. Run `from huggingface_hub import login; login()` in Colab first.")
    api = HfApi()
    username = api.whoami()["name"]
    if not args.repo_id.startswith(username + "/"):
        raise ValueError(f"Logged in as {username}, but repo-id is {args.repo_id}")

    output_dir = Path("hub_upload") / run_name
    output_dir.mkdir(parents=True, exist_ok=True)
    import shutil
    shutil.copy2(model_path, output_dir / "model.pt")
    shutil.copy2(__file__, output_dir / "ppo.py")
    scores = []
    frames = []
    for episode in range(args.eval_episodes):
        env = gym.make(args.env_id, render_mode="rgb_array")
        obs, _ = env.reset(seed=args.seed + 10_000 + episode)
        done = False
        total_reward = 0.0
        step_count = 0
        while not done:
            x = torch.as_tensor(obs, dtype=torch.float32, device=device).unsqueeze(0)
            with torch.no_grad():
                action = agent.actor(x).argmax(dim=-1).item()
            obs, reward, terminated, truncated, _ = env.step(action)
            total_reward += float(reward)
            done = terminated or truncated
            if episode == 0:
                frame = env.render()
                if frame is not None:
                    frames.append(frame)
            step_count += 1
            if step_count >= 2000:
                break
        env.close()
        scores.append(total_reward)
        print(f"Evaluation episode {episode + 1}/{args.eval_episodes}: reward={total_reward:.1f}")

    if frames:
        video_path = output_dir / "replay.mp4"
        with imageio.get_writer(video_path, fps=30, codec="libx264", quality=7) as writer:
            for frame in frames:
                writer.append_data(frame)

    mean_reward = float(np.mean(scores))
    std_reward = float(np.std(scores))
    (output_dir / "eval_results.json").write_text(
        json.dumps({"env_id": args.env_id, "mean_reward": mean_reward,
                    "std_reward": std_reward, "episode_rewards": scores}, indent=2),
        encoding="utf-8",
    )
    (output_dir / "hyperparameters.json").write_text(json.dumps(vars(args), indent=2), encoding="utf-8")
    (output_dir / "README.md").write_text(
        f"""---
license: mit
tags:
- reinforcement-learning
- ppo
- gymnasium
- deep-rl-course
model-index:
- name: PPO {args.env_id}
  results:
  - task:
      type: reinforcement-learning
      name: Reinforcement Learning
    dataset:
      name: {args.env_id}
      type: {args.env_id}
    metrics:
    - type: mean_reward
      value: {mean_reward:.3f}
      name: Mean Reward
---
# PPO from scratch — {args.env_id}

Trained using a CleanRL-style PyTorch PPO implementation adapted for Gymnasium.

Evaluation over {args.eval_episodes} episodes: **{mean_reward:.2f} ± {std_reward:.2f}**.

- `model.pt`: actor-critic PyTorch state dict
- `ppo.py`: full training/evaluation/upload source
- `replay.mp4`: evaluation video (if rendering succeeded)
- `hyperparameters.json`: training settings
- `eval_results.json`: evaluation results

To run locally, install PyTorch, `gymnasium[box2d]`, `imageio`, `imageio-ffmpeg`, `tensorboard`, and `huggingface_hub`.
""",
        encoding="utf-8",
    )
    print(f"Uploading to https://huggingface.co/{args.repo_id}")
    api.create_repo(repo_id=args.repo_id, repo_type="model", exist_ok=True)
    api.upload_folder(folder_path=str(output_dir), repo_id=args.repo_id, repo_type="model",
                      commit_message="Upload PPO model, evaluation and replay")
    print(f"Upload complete: https://huggingface.co/{args.repo_id}")


def main():
    args = parse_args()
    run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = True
    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
    writer = SummaryWriter(f"runs/{run_name}")
    writer.add_text("hyperparameters", "```json\n" + json.dumps(vars(args), indent=2) + "\n```")

    envs = gym.vector.SyncVectorEnv([make_env(args.env_id, args.seed + i) for i in range(args.num_envs)])
    if not isinstance(envs.single_action_space, gym.spaces.Discrete):
        raise ValueError("This PPO script supports discrete action spaces only")
    agent = Agent(envs).to(device)
    optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)
    obs_shape = envs.single_observation_space.shape
    obs = torch.zeros((args.num_steps, args.num_envs) + obs_shape, device=device)
    actions = torch.zeros((args.num_steps, args.num_envs), device=device)
    logprobs = torch.zeros((args.num_steps, args.num_envs), device=device)
    rewards = torch.zeros((args.num_steps, args.num_envs), device=device)
    dones = torch.zeros((args.num_steps, args.num_envs), device=device)
    values = torch.zeros((args.num_steps, args.num_envs), device=device)
    global_step = 0
    start_time = time.time()
    next_obs_np, _ = envs.reset(seed=args.seed)
    next_obs = torch.as_tensor(next_obs_np, dtype=torch.float32, device=device)
    next_done = torch.zeros(args.num_envs, device=device)
    num_updates = max(1, args.total_timesteps // args.batch_size)

    for update in range(1, num_updates + 1):
        if args.anneal_lr:
            frac = 1.0 - (update - 1.0) / num_updates
            optimizer.param_groups[0]["lr"] = frac * args.learning_rate
        for step in range(args.num_steps):
            global_step += args.num_envs
            obs[step] = next_obs
            dones[step] = next_done
            with torch.no_grad():
                action, logprob, _, value = agent.get_action_and_value(next_obs)
                values[step] = value.flatten()
            actions[step] = action
            logprobs[step] = logprob
            next_obs_np, reward, terminated, truncated, infos = envs.step(action.cpu().numpy())
            done = np.logical_or(terminated, truncated)
            rewards[step] = torch.as_tensor(reward, dtype=torch.float32, device=device)
            next_obs = torch.as_tensor(next_obs_np, dtype=torch.float32, device=device)
            next_done = torch.as_tensor(done, dtype=torch.float32, device=device)
            for info in (infos.get("final_info") if infos.get("final_info") is not None else []):
                if info and "episode" in info:
                    episode = info["episode"]
                    ep_r = float(np.asarray(episode["r"]).item())
                    ep_l = int(np.asarray(episode["l"]).item())
                    print(f"global_step={global_step}, episodic_return={ep_r:.2f}")
                    writer.add_scalar("charts/episodic_return", ep_r, global_step)
                    writer.add_scalar("charts/episodic_length", ep_l, global_step)
            # Gymnasium >= 1.0 can report vectorized episode data directly.
            if "episode" in infos and "_episode" in infos:
                for i, valid in enumerate(infos["_episode"]):
                    if valid:
                        writer.add_scalar("charts/episodic_return", float(infos["episode"]["r"][i]), global_step)

        with torch.no_grad():
            next_value = agent.get_value(next_obs).reshape(1, -1)
            advantages = torch.zeros_like(rewards)
            lastgaelam = 0
            for t in reversed(range(args.num_steps)):
                if t == args.num_steps - 1:
                    nextnonterminal = 1.0 - next_done
                    nextvalues = next_value
                else:
                    nextnonterminal = 1.0 - dones[t + 1]
                    nextvalues = values[t + 1]
                delta = rewards[t] + args.gamma * nextvalues * nextnonterminal - values[t]
                advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam
            returns = advantages + values

        b_obs = obs.reshape((-1,) + obs_shape)
        b_logprobs = logprobs.reshape(-1)
        b_actions = actions.reshape(-1)
        b_advantages = advantages.reshape(-1)
        b_returns = returns.reshape(-1)
        b_values = values.reshape(-1)
        b_inds = np.arange(args.batch_size)
        clipfracs = []
        for epoch in range(args.update_epochs):
            np.random.shuffle(b_inds)
            for start in range(0, args.batch_size, args.minibatch_size):
                mb_inds = b_inds[start:start + args.minibatch_size]
                _, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions.long()[mb_inds])
                logratio = newlogprob - b_logprobs[mb_inds]
                ratio = logratio.exp()
                with torch.no_grad():
                    old_approx_kl = (-logratio).mean()
                    approx_kl = ((ratio - 1) - logratio).mean()
                    clipfracs.append(((ratio - 1.0).abs() > args.clip_coef).float().mean().item())
                mb_advantages = b_advantages[mb_inds]
                if args.norm_adv:
                    mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
                pg_loss1 = -mb_advantages * ratio
                pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
                pg_loss = torch.max(pg_loss1, pg_loss2).mean()
                newvalue = newvalue.view(-1)
                if args.clip_vloss:
                    v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2
                    v_clipped = b_values[mb_inds] + torch.clamp(newvalue - b_values[mb_inds], -args.clip_coef, args.clip_coef)
                    v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2
                    v_loss = 0.5 * torch.max(v_loss_unclipped, v_loss_clipped).mean()
                else:
                    v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()
                entropy_loss = entropy.mean()
                loss = pg_loss - args.ent_coef * entropy_loss + args.vf_coef * v_loss
                optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)
                optimizer.step()
            if args.target_kl is not None and approx_kl > args.target_kl:
                break
        y_pred, y_true = b_values.cpu().numpy(), b_returns.cpu().numpy()
        var_y = np.var(y_true)
        explained_var = np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
        writer.add_scalar("charts/learning_rate", optimizer.param_groups[0]["lr"], global_step)
        writer.add_scalar("losses/value_loss", v_loss.item(), global_step)
        writer.add_scalar("losses/policy_loss", pg_loss.item(), global_step)
        writer.add_scalar("losses/entropy", entropy_loss.item(), global_step)
        writer.add_scalar("losses/old_approx_kl", old_approx_kl.item(), global_step)
        writer.add_scalar("losses/approx_kl", approx_kl.item(), global_step)
        writer.add_scalar("losses/clipfrac", np.mean(clipfracs), global_step)
        writer.add_scalar("losses/explained_variance", explained_var, global_step)
        sps = int(global_step / max(time.time() - start_time, 1e-6))
        writer.add_scalar("charts/SPS", sps, global_step)
        print(f"Update {update}/{num_updates} | Steps {global_step} | SPS {sps} | Policy loss {pg_loss.item():.4f}")

    output_dir = Path("models")
    output_dir.mkdir(exist_ok=True)
    model_path = output_dir / f"{run_name}.pt"
    torch.save(agent.state_dict(), model_path)
    print(f"Training completed! Model saved to {model_path}")
    envs.close()
    writer.close()
    if args.upload:
        evaluate_and_upload(agent, args, device, model_path, run_name)
    else:
        print("To evaluate and upload, rerun with --upload (this starts a new training run).")


if __name__ == "__main__":
    main()