"""Chapter 9 lab: Proximal Policy Optimization (PPO-clip) on CartPole.

Your tasks (in order; run `python -m pytest -q` after each):

    Task 1  clipped_surrogate_loss   the PPO-clip objective, as a loss to minimize
    Task 2  approx_kl, clip_fraction the two numbers you watch while PPO trains
    Task 3  value_loss               the critic's loss, with optional value clipping
    Task 4  ppo_update               K epochs of minibatch updates on one batch

`train_ppo` at the bottom is complete: it alternates collect_rollout (common.py)
and your ppo_update. When all tests pass, run `python run.py`.

Shapes: every per-sample tensor below has shape (B,), where B is the (mini)batch size.
"""

import time
from dataclasses import dataclass

import numpy as np
import torch

from cartpole import CartPole
from common import ActorCritic, collect_rollout


# ---------------------------------------------------------------------------
# Task 1: the clipped surrogate loss
# ---------------------------------------------------------------------------
def clipped_surrogate_loss(logp_new, logp_old, adv, clip_eps):
    """PPO-clip loss (the chapter's "clipped surrogate objective", with a minus sign).

    Inputs (all shape (B,)):
      logp_new: log pi_theta(a_t | s_t) under the current parameters (requires grad)
      logp_old: log pi_old(a_t | s_t), recorded when the data was sampled (no grad)
      adv:      advantage estimates A_hat_t
      clip_eps: epsilon, e.g. 0.2

    Returns a scalar tensor:

        rho_t = pi_theta(a_t|s_t) / pi_old(a_t|s_t) = exp(logp_new - logp_old)
        loss  = - mean_t [ min( rho_t * A_t,  clip(rho_t, 1 - eps, 1 + eps) * A_t ) ]

    The minus sign is there because PyTorch optimizers minimize, and we want to
    MAXIMIZE the clipped objective.
    Hint: torch.exp, torch.clamp, torch.min (elementwise min of two tensors).
    """
    # TODO: compute the ratio, the unclipped and clipped terms, and return the loss.
    raise NotImplementedError("Task 1: clipped_surrogate_loss")


# ---------------------------------------------------------------------------
# Task 2: diagnostics
# ---------------------------------------------------------------------------
@torch.no_grad()
def approx_kl(logp_new, logp_old):
    """Estimate KL(pi_old || pi_theta) from samples a_t ~ pi_old. Returns a Python float.

    With rho_t = exp(logp_new - logp_old), use the low-variance, always-nonnegative
    estimator from the chapter (called k3 in Chapter 11):

        approx_kl = mean_t [ (rho_t - 1) - log rho_t ]

    It is 0 when the policies agree on the sampled actions.
    """
    # TODO
    raise NotImplementedError("Task 2: approx_kl")


@torch.no_grad()
def clip_fraction(logp_new, logp_old, clip_eps):
    """Fraction of samples whose ratio lies outside [1 - eps, 1 + eps]. Returns a Python float.

        clip_frac = mean_t [ |rho_t - 1| > eps ]

    This counts samples that are outside the band, whichever sign their advantage has.
    """
    # TODO
    raise NotImplementedError("Task 2: clip_fraction")


# ---------------------------------------------------------------------------
# Task 3: the value loss (optionally clipped)
# ---------------------------------------------------------------------------
def value_loss(values, returns, values_old=None, clip_eps=None):
    """Critic loss. All tensors shape (B,). Returns a scalar tensor.

    Without clipping (clip_eps is None):

        L_V = 0.5 * mean_t (V(s_t) - R_t)^2

    With clipping (clip_eps is a number; values_old = V_old(s_t) from the rollout):

        V_clip = V_old + clip(V - V_old, -eps, +eps)
        L_V    = 0.5 * mean_t max( (V - R_t)^2, (V_clip - R_t)^2 )

    (Value clipping is a PPO implementation detail: the critic, like the policy,
    is discouraged from moving far from its old predictions in one batch.)
    """
    # TODO
    raise NotImplementedError("Task 3: value_loss")


# ---------------------------------------------------------------------------
# Task 4: the PPO update
# ---------------------------------------------------------------------------
@dataclass
class PPOConfig:
    n_steps: int = 512          # environment steps per rollout (the "batch")
    epochs: int = 10            # K: passes over each batch
    minibatch_size: int = 64
    clip_eps: float = 0.2       # epsilon in the clipped surrogate
    lr: float = 1e-3            # Adam learning rate (annealed linearly to 0 if anneal_lr)
    gamma: float = 0.99
    lam: float = 0.95
    vf_coef: float = 0.5        # c1: weight of the value loss
    ent_coef: float = 0.0       # c2: weight of the entropy bonus
    max_grad_norm: float = 0.5  # global gradient-norm clip
    clip_value: bool = False    # use value clipping in the critic loss
    norm_adv: bool = True       # normalize advantages per minibatch
    target_kl: float = None     # stop the epochs early if approx KL exceeds this
    anneal_lr: bool = True


def ppo_update(agent, optimizer, batch, cfg, generator=None):
    """Run up to cfg.epochs passes over `batch`, in shuffled minibatches. Returns a stats dict.

    agent:     ActorCritic (common.py). agent.dist(obs) -> Categorical; agent.value(obs) -> (B,)
    optimizer: a torch optimizer over agent.parameters()
    batch:     dict from collect_rollout with keys
               obs (N,4), actions (N,), logp_old (N,), values_old (N,), advantages (N,), returns (N,)
    cfg:       PPOConfig
    generator: torch.Generator for the shuffles (pass it to torch.randperm)

    For each epoch (up to cfg.epochs):
      Shuffle the indices 0..N-1 (torch.randperm(N, generator=generator)) and cut them into
      minibatches of cfg.minibatch_size. For each minibatch:
        1. Evaluate the CURRENT policy on the minibatch states: the log-probabilities of the
           stored actions (logp_new, with gradient) and the mean entropy of pi_theta(.|s).
        2. If cfg.norm_adv, normalize the minibatch advantages to mean 0 and std 1
           (add 1e-8 to the std so you never divide by zero).
        3. Build the loss from your Tasks 1 and 3:
              policy loss  pl = clipped surrogate loss (uses the STORED logp_old, never recomputed)
              value loss   vl = value loss of V_phi(s) against the stored returns; pass
                                values_old and cfg.clip_eps only if cfg.clip_value, else None
              loss = pl + cfg.vf_coef * vl - cfg.ent_coef * entropy
        4. One optimizer step: zero the gradients, backpropagate, clip the global gradient
           norm to cfg.max_grad_norm, step.
        5. Record pl, vl, entropy (as floats) and your Task 2 diagnostics approx_kl and
           clip_fraction, computed from this minibatch's logp_new and logp_old.
      After the epoch: if cfg.target_kl is not None, compute approx_kl on the FULL batch with
      the current policy (under torch.no_grad()); if it exceeds cfg.target_kl, stop the epochs.

    Useful calls: agent.dist(obs) -> Categorical, with .log_prob(actions) and .entropy();
    agent.value(obs) -> (B,); torch.nn.utils.clip_grad_norm_(agent.parameters(), max_norm).

    Return a dict with the MEAN over all minibatch steps of
      "policy_loss", "value_loss", "entropy", "approx_kl", "clip_frac"   (Python floats)
    plus "epochs_run": the number of epochs completed (an int; 1 <= epochs_run <= cfg.epochs).
    """
    # TODO: implement the loop described above.
    raise NotImplementedError("Task 4: ppo_update")


# ---------------------------------------------------------------------------
# The training loop (given, nothing to do here)
# ---------------------------------------------------------------------------
def train_ppo(cfg=None, total_steps=50_000, seed=0, log_every=0):
    """Train PPO on CartPole. Returns a history dict of per-update lists:

        steps, mean_return (mean of the last 20 finished episodes), approx_kl, clip_frac,
        entropy, value_loss, epochs_run
    """
    cfg = cfg or PPOConfig()
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed)
    env = CartPole(seed=seed)
    agent = ActorCritic()
    optimizer = torch.optim.Adam(agent.parameters(), lr=cfg.lr, eps=1e-5)

    obs = env.reset()
    n_updates = total_steps // cfg.n_steps
    recent = []
    hist = {k: [] for k in ["steps", "mean_return", "approx_kl", "clip_frac", "entropy", "value_loss", "epochs_run"]}
    t0 = time.time()
    for u in range(n_updates):
        if cfg.anneal_lr:
            optimizer.param_groups[0]["lr"] = cfg.lr * (1.0 - u / n_updates)
        batch, obs, ep_returns = collect_rollout(env, agent, obs, cfg.n_steps, cfg.gamma, cfg.lam)
        recent = (recent + ep_returns)[-20:]
        stats = ppo_update(agent, optimizer, batch, cfg, generator=gen)

        hist["steps"].append((u + 1) * cfg.n_steps)
        hist["mean_return"].append(float(np.mean(recent)) if recent else 0.0)
        for k in ["approx_kl", "clip_frac", "entropy", "value_loss", "epochs_run"]:
            hist[k].append(stats[k])
        if log_every and (u + 1) % log_every == 0:
            print(f"step {hist['steps'][-1]:6d}  return {hist['mean_return'][-1]:6.1f}  "
                  f"kl {stats['approx_kl']:.4f}  clipfrac {stats['clip_frac']:.2f}  "
                  f"entropy {stats['entropy']:.3f}  ({time.time() - t0:.0f}s)")
    hist["agent"] = agent
    return hist