"""Chapter 18 lab: a multi-turn tool-using agent trained with group-relative policy gradient.

Your file: fill in the four TODOs. (The reference solution is solutions/agentic.py.)

The task. A secret number s in 0..7. Half the episodes reveal s in the prompt ("easy");
the other half hide it ("hard"). The agent may call a tool, compare(x), which answers
"lower" (s < x), "higher" (s > x) or "equal" (s == x). Each call costs a little reward.
The episode ends when the agent answers, writes something malformed, or runs out of turns.

    easy:  [<bos> 5]  ans 5                                         reward 1
    hard:  [<bos> ?]  call 3 {<tool> higher}  call 5 {<tool> lower}  ans 4
                                                     turn rewards -0.05, -0.05, +1

What you build (four tasks, in order):
    Task 1  GuessEnv.step            the environment: run the tool, hand out the reward
    Task 2  rollout                  interleave sampled model tokens with tool tokens
    Task 3  turn_advantages          trajectory-level vs turn-level credit, group baseline
    Task 4  masked_policy_loss       REINFORCE on model tokens only
"""

import numpy as np
import torch
import torch.nn.functional as F

from tinyagent import (
    ANS, BOS, CALL, EQUAL, HIDDEN, HIGHER, LOWER, N_NUMBERS, TOOL, TURN_LEN,
    Policy, Trajectory, is_number, make_batch, sample_token, warm_start,
)

torch.set_num_threads(1)


# ---------------------------------------------------------------------------
# Task 1: the environment
# ---------------------------------------------------------------------------
class GuessEnv:
    """Find a secret number in 0..N_NUMBERS-1, optionally using a compare(x) tool.

    Args:
        secret:    the secret number s.
        revealed:  if True the prompt shows s ("easy"); if False it shows "?" ("hard").
        call_cost: reward subtracted for every tool call (a stand-in for latency and cost).
        max_turns: the episode is cut off after this many model turns.
        malformed_reward: reward for a turn that is not a valid "call x" or "ans x".
                   It must be clearly negative: see the chapter's exercise on why 0 is a bug.
    """

    def __init__(self, secret, revealed, call_cost=0.05, max_turns=5, malformed_reward=-1.0):
        self.secret, self.revealed = int(secret), bool(revealed)
        self.call_cost, self.max_turns = call_cost, max_turns
        self.malformed_reward = malformed_reward
        self.turns = 0

    def reset(self):
        """Start an episode. Returns the prompt tokens. (Provided.)"""
        self.turns = 0
        return [BOS, self.secret if self.revealed else HIDDEN]

    def step(self, turn):
        """Task 1. Process one model turn and return (obs_tokens, reward, done).

        Args:
            turn: a list of TURN_LEN = 2 token ids the model just wrote: [verb, number].

        Rules (the chapter's toy environment):
            * Count the turn (self.turns += 1).
            * ["call", x] with x a number: the tool compares x with the secret and returns
              obs_tokens = [TOOL, LOWER] if secret < x, [TOOL, HIGHER] if secret > x,
              [TOOL, EQUAL] if secret == x. reward = -self.call_cost. done = False,
              unless this was turn number self.max_turns (then done = True: out of turns).
            * ["ans", x] with x a number: obs_tokens = [], reward = 1.0 if x == secret else 0.0,
              done = True.
            * Anything else is a malformed turn: obs_tokens = [], reward = self.malformed_reward,
              done = True (like a tool call that fails to parse: the episode is halted).

        Use is_number(tok) to check that a token is one of the digits 0..7.
        """
        # TODO: count the turn, unpack verb and number, and return (obs_tokens, reward, done)
        #       for the three cases: a valid call, a valid answer, and anything else.
        raise NotImplementedError("Task 1: GuessEnv.step")

# ---------------------------------------------------------------------------
# Task 2: the rollout loop
# ---------------------------------------------------------------------------
def rollout(sampler, env, rng):
    """Task 2. Play one episode and record who wrote each token.

    Args:
        sampler: a tinyagent.Sampler. sampler.start() gives a fresh state;
                 probs, state = sampler.step(state, token) feeds one token and returns the
                 next-token probabilities.
        env:     a GuessEnv.
        rng:     a numpy Generator; draw tokens with sample_token(probs, rng).

    Returns:
        a tinyagent.Trajectory with tokens, is_model, turn_id and turn_rewards filled in.

    The loop:
        1. tokens = env.reset(). Record them with is_model = 0, turn_id = -1, and feed each one
           to the sampler (the policy must read the prompt).
        2. For turn k = 0, 1, 2, ...:
             a. sample TURN_LEN tokens one at a time from the current probs; record each with
                is_model = 1, turn_id = k, and feed it to the sampler;
             b. obs, r, done = env.step(those TURN_LEN tokens); append r to turn_rewards;
             c. record each obs token with is_model = 0, turn_id = k, and feed it to the
                sampler (the policy reads the tool output, it does not write it);
             d. stop when done.
    """
    # TODO: follow steps 1-2 of the docstring. Keep three parallel lists in a Trajectory
    #       (tokens, is_model, turn_id) and one reward per turn. Remember to feed EVERY token
    #       to the sampler, including the prompt and the tool's output.
    raise NotImplementedError("Task 2: rollout")

# ---------------------------------------------------------------------------
# Task 3: credit assignment across turns
# ---------------------------------------------------------------------------
def turn_advantages(group_rewards, mode="trajectory", gamma=1.0):
    """Task 3. Per-turn advantages for a group of G trajectories that share one task.

    Args:
        group_rewards: list of G lists; group_rewards[i][k] is the reward of turn k of
                       trajectory i. Trajectories may have different numbers of turns.
        mode:          "trajectory" or "turn".
        gamma:         discount across turns (used only in "turn" mode).

    Returns:
        list of G float numpy arrays; out[i][k] is the advantage of turn k of trajectory i.

    "trajectory" (the GRPO-style default; no std division, as in Dr. GRPO):
        R_i = sum_k r_{i,k};   b = mean_i R_i;   every turn of trajectory i gets  R_i - b.

    "turn" (reward-to-go with a group baseline):
        G_{i,k} = sum_{j >= k} gamma^(j - k) r_{i,j}      (the return from turn k on)
        b = mean_i G_{i,0};   turn k of trajectory i gets  G_{i,k} - b.

    With gamma = 1 the first turn gets the same advantage in both modes; later turns
    differ in "turn" mode because they are not charged for rewards earned before them.
    """
    # TODO: implement both modes. For "turn", compute the reward-to-go backwards
    #       (g_k = r_k + gamma * g_{k+1}), then subtract the group baseline.
    raise NotImplementedError("Task 3: turn_advantages")

# ---------------------------------------------------------------------------
# Task 4: the masked policy-gradient loss
# ---------------------------------------------------------------------------
def masked_policy_loss(logits, ids, loss_mask, token_adv):
    """Task 4. The REINFORCE surrogate loss over the tokens selected by loss_mask.

    Args:
        logits:    (B, T-1, V) = policy(ids[:, :-1]); position t predicts token t+1.
        ids:       (B, T) token ids of the padded trajectories.
        loss_mask: (B, T) 1.0 for tokens that should receive gradient (the model's own tokens).
        token_adv: (B, T) the advantage attached to each token.

    Returns:
        a scalar tensor
            loss = - sum_{b,t} m_{b,t} * A_{b,t} * log pi(ids_{b,t} | ids_{b,<t})  /  sum_{b,t} m_{b,t}
        where the sum runs over t = 1..T-1 (token 0 has no prediction) and m = loss_mask.
        (This is the "token-mean" aggregation of Chapter 16.) Guard against an all-zero mask.

    Steps: log_softmax the logits, gather the log-prob of the token that actually came next
    (ids[:, 1:]), and line up loss_mask and token_adv by dropping their first column too.
    """
    # TODO: log_softmax, gather the next-token log-probs, align the mask and advantages
    #       (drop column 0), and return the masked token-mean of -A * log pi.
    raise NotImplementedError("Task 4: masked_policy_loss")

# ---------------------------------------------------------------------------
# The training loop (provided, complete)
# ---------------------------------------------------------------------------
def train(iters=250, tasks_per_iter=32, group_size=8, mode="trajectory", gamma=1.0,
          mask_env=True, call_cost=0.05, lr=2e-3, seed=0, warm_steps=500, log_every=0):
    """Group-relative policy gradient on GuessEnv. Returns (policy, history).

    Each iteration: sample tasks_per_iter tasks (secret, revealed); play group_size episodes
    of each with a frozen snapshot of the policy; compute per-turn advantages within each
    group (Task 3); spread them over the tokens; take one Adam step on the masked loss
    (Task 4). history is a list of dicts with the metrics of every iteration.
    """
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    policy = warm_start(Policy(), steps=warm_steps, seed=seed)
    opt = torch.optim.Adam(policy.parameters(), lr=lr)
    history = []
    for it in range(iters):
        sampler = policy.sampler()                      # "inference engine" snapshot
        trajs, advs, revealed = [], [], []
        for _ in range(tasks_per_iter):
            secret, rev = int(rng.integers(N_NUMBERS)), bool(rng.random() < 0.5)
            group = [rollout(sampler, GuessEnv(secret, rev, call_cost=call_cost), rng)
                     for _ in range(group_size)]
            advs += turn_advantages([t.turn_rewards for t in group], mode=mode, gamma=gamma)
            trajs += group
            revealed += [rev] * group_size
        ids, loss_mask, token_adv = make_batch(trajs, advs, mask_env=mask_env)
        loss = masked_policy_loss(policy(ids[:, :-1]), ids, loss_mask, token_adv)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
        opt.step()

        rev = np.array(revealed)
        solved = np.array([t.solved for t in trajs])
        calls = np.array([t.n_calls for t in trajs])
        n_model = sum(sum(t.is_model) for t in trajs)
        n_tool = sum(len(t.tokens) - sum(t.is_model) - 2 for t in trajs)
        history.append({
            "iter": it,
            "reward": float(np.mean([t.total_reward for t in trajs])),
            "solved_easy": float(solved[rev].mean()) if rev.any() else float("nan"),
            "solved_hard": float(solved[~rev].mean()) if (~rev).any() else float("nan"),
            "calls_easy": float(calls[rev].mean()) if rev.any() else float("nan"),
            "calls_hard": float(calls[~rev].mean()) if (~rev).any() else float("nan"),
            "tool_token_frac": n_tool / max(1, n_tool + n_model),
        })
        if log_every and it % log_every == 0:
            h = history[-1]
            print(f"iter {it:4d}  reward {h['reward']:+.3f}  solved easy {h['solved_easy']:.2f} "
                  f"hard {h['solved_hard']:.2f}  calls easy {h['calls_easy']:.2f} hard {h['calls_hard']:.2f}")
    return policy, history