"""Helper code for the PPO lab. You do NOT need to edit this file.

It contains the parts of PPO that are not new in this chapter:

    ActorCritic      a small policy network + value network for CartPole
    collect_rollout  run the current policy for n_steps and record everything PPO needs
    compute_gae      generalized advantage estimation (you wrote this in Chapter 8's lab)

Read `collect_rollout` once: the arrays it returns are exactly the `batch` that
`ppo_update` in ppo.py receives.
"""

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

torch.set_num_threads(1)  # tiny networks are faster on one thread


def _layer(n_in, n_out, gain):
    """Linear layer with orthogonal init (a common PPO implementation detail)."""
    layer = nn.Linear(n_in, n_out)
    nn.init.orthogonal_(layer.weight, gain)
    nn.init.zeros_(layer.bias)
    return layer


class ActorCritic(nn.Module):
    """Separate actor (policy) and critic (value) MLPs, 2 hidden layers of tanh units.

    The policy's last layer starts with tiny weights (gain 0.01) so the initial policy
    is close to uniform: pi(left) ~ pi(right) ~ 0.5 in every state.
    """

    def __init__(self, obs_dim=4, n_actions=2, hidden=64):
        super().__init__()
        g = np.sqrt(2)
        self.actor = nn.Sequential(
            _layer(obs_dim, hidden, g), nn.Tanh(),
            _layer(hidden, hidden, g), nn.Tanh(),
            _layer(hidden, n_actions, 0.01),
        )
        self.critic = nn.Sequential(
            _layer(obs_dim, hidden, g), nn.Tanh(),
            _layer(hidden, hidden, g), nn.Tanh(),
            _layer(hidden, 1, 1.0),
        )

    def dist(self, obs):
        """Categorical distribution pi_theta(. | s) for a batch of observations, shape (B, obs_dim)."""
        return Categorical(logits=self.actor(obs))

    def value(self, obs):
        """V_phi(s), shape (B,)."""
        return self.critic(obs).squeeze(-1)


def compute_gae(rewards, values, dones, last_value, gamma=0.99, lam=0.95):
    """Generalized advantage estimation (Chapter 8).

    rewards, values, dones: arrays of shape (T,). dones[t] = 1.0 if the episode ended
    after step t (so V(s_{t+1}) must not be bootstrapped from).
    last_value: V(s_T), the value of the state after the final step.

        delta_t = r_t + gamma * V(s_{t+1}) * (1 - done_t) - V(s_t)
        A_t     = delta_t + gamma * lam * (1 - done_t) * A_{t+1}

    Returns (advantages, returns) with returns = advantages + values (the critic's target).
    """
    T = len(rewards)
    adv = np.zeros(T, dtype=np.float32)
    next_value = last_value
    running = 0.0
    for t in reversed(range(T)):
        nonterminal = 1.0 - dones[t]
        delta = rewards[t] + gamma * next_value * nonterminal - values[t]
        running = delta + gamma * lam * nonterminal * running
        adv[t] = running
        next_value = values[t]
    return adv, adv + np.asarray(values, dtype=np.float32)


@torch.no_grad()
def collect_rollout(env, agent, obs, n_steps, gamma=0.99, lam=0.95):
    """Run the current policy for n_steps (across episode boundaries).

    obs: the current observation of `env` (continue from where the last rollout stopped).

    Returns (batch, obs, episode_returns):
      batch: dict of torch tensors, each with first dimension n_steps
         "obs"        (T, 4)  states s_t
         "actions"    (T,)    actions a_t (int64)
         "logp_old"   (T,)    log pi_old(a_t | s_t), recorded at sampling time
         "values_old" (T,)    V_old(s_t), recorded at sampling time
         "advantages" (T,)    GAE advantages A_hat_t (NOT normalized)
         "returns"    (T,)    targets for the critic: A_hat_t + V_old(s_t)
      obs: the observation to continue from next time
      episode_returns: list of the total rewards of episodes that finished in this rollout

    Note: CartPole's 500-step time limit is treated as a true terminal state. That is a
    small simplification (strictly, a time-out should bootstrap from V); it does not
    matter for this lab.
    """
    obs_buf = np.zeros((n_steps, len(obs)), dtype=np.float32)
    act_buf = np.zeros(n_steps, dtype=np.int64)
    logp_buf = np.zeros(n_steps, dtype=np.float32)
    val_buf = np.zeros(n_steps, dtype=np.float32)
    rew_buf = np.zeros(n_steps, dtype=np.float32)
    done_buf = np.zeros(n_steps, dtype=np.float32)
    episode_returns = []
    ep_ret = getattr(env, "_ep_ret", 0.0)

    for t in range(n_steps):
        o = torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0)
        d = agent.dist(o)
        a = d.sample()
        obs_buf[t] = obs
        act_buf[t] = a.item()
        logp_buf[t] = d.log_prob(a).item()
        val_buf[t] = agent.value(o).item()

        obs, r, done = env.step(a.item())
        rew_buf[t] = r
        done_buf[t] = float(done)
        ep_ret += r
        if done:
            episode_returns.append(ep_ret)
            ep_ret = 0.0
            obs = env.reset()
    env._ep_ret = ep_ret

    last_value = agent.value(torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0)).item()
    adv, ret = compute_gae(rew_buf, val_buf, done_buf, last_value, gamma, lam)
    batch = {
        "obs": torch.as_tensor(obs_buf),
        "actions": torch.as_tensor(act_buf),
        "logp_old": torch.as_tensor(logp_buf),
        "values_old": torch.as_tensor(val_buf),
        "advantages": torch.as_tensor(adv),
        "returns": torch.as_tensor(ret),
    }
    return batch, obs, episode_returns