"""Chapter 20 lab, Task 1: where the GPU memory goes.

Three small functions. Run `python -m pytest -q -k task1` after each.

    Task 1a  model_state_bytes   bytes to hold one model, trained or frozen
    Task 1b  kv_cache_bytes      bytes of KV cache for some number of tokens
    Task 1c  rl_memory_gb        the memory bill of a PPO or GRPO setup, by component

Conventions (the chapter's "Memory accounting" section):
  * A TRAINED model with mixed-precision Adam costs 16 bytes per parameter:
        bf16 weights (2) + bf16 gradients (2) + fp32 master weights (4)
        + Adam first moment m (4) + Adam second moment v (4) = 16.
  * A FROZEN model (reference, reward model) costs only its weights:
        2 bytes per parameter in bf16.
  * We ignore activations, temporary buffers and fragmentation. Real runs need
    noticeably more than these numbers; treat them as a floor.
  * GB means 1e9 bytes (the unit on GPU spec sheets).
"""

GB = 1e9
TRAIN_BYTES_PER_PARAM = 16


# ---------------------------------------------------------------------------
# Task 1a
# ---------------------------------------------------------------------------
def model_state_bytes(n_params, trainable, bytes_per_weight=2):
    """Bytes needed to hold one model.

    n_params:          number of parameters (e.g. 7e9)
    trainable:         True if we train it with mixed-precision Adam, False if frozen
    bytes_per_weight:  bytes per stored weight for a FROZEN model (2 for bf16)

    Returns TRAIN_BYTES_PER_PARAM * n_params if trainable, else bytes_per_weight * n_params.
    Example: a trained 7B model needs 16 * 7e9 = 1.12e11 bytes = 112 GB.
    """
    # TODO
    raise NotImplementedError("Task 1a: model_state_bytes")


# ---------------------------------------------------------------------------
# Task 1b
# ---------------------------------------------------------------------------
def kv_cache_bytes(n_layers, n_kv_heads, head_dim, n_tokens, bytes_per_value=2):
    """Bytes of KV cache for `n_tokens` tokens (summed over all sequences).

    For every token, every layer stores one KEY vector and one VALUE vector for each
    KV head, each of length head_dim:

        bytes = 2 * n_layers * n_kv_heads * head_dim * n_tokens * bytes_per_value

    Example (the chapter's worked example, Qwen2.5-7B's shape: 28 layers, 4 KV heads,
    head_dim 128, bf16): 2 * 28 * 4 * 128 * 2 = 57,344 bytes = 56 KiB per token.
    """
    # TODO
    raise NotImplementedError("Task 1b: kv_cache_bytes")


# ---------------------------------------------------------------------------
# Task 1c
# ---------------------------------------------------------------------------
def rl_memory_gb(algo, n_params, critic_params=None, rm_params=None, use_ref=True, lora_params=0):
    """Memory (in GB) for the model states of one RL setup, by component.

    algo:          'ppo' or 'grpo' (raise ValueError otherwise)
    n_params:      parameters of the policy
    critic_params: parameters of the critic (PPO only); default: same as the policy
    rm_params:     parameters of a learned reward model. PPO always has one
                   (default: same size as the policy). GRPO has one only if
                   rm_params is given; otherwise it uses a verifier, which costs 0.
    use_ref:       keep a frozen reference model for the KL term? (False when beta = 0)
    lora_params:   if > 0, train a LoRA adapter with this many parameters instead of
                   the full policy.

    Return a dict with keys 'policy', 'reference', 'critic', 'reward_model', 'total'
    (all floats, in GB = 1e9 bytes):

      policy:       full fine-tuning: trainable n_params.
                    LoRA: frozen n_params (bf16) + trainable lora_params.
      reference:    frozen n_params if use_ref, else 0.
                    With LoRA it is 0 even if use_ref: the reference model IS the base
                    model with the adapter switched off, so no extra copy is needed.
      critic:       PPO: trainable critic_params (default n_params). GRPO: 0.
      reward_model: PPO: frozen rm_params (default n_params).
                    GRPO: frozen rm_params if given, else 0.
      total:        the sum of the four.

    Example: rl_memory_gb('ppo', 7e9)['total'] == 252.0 (112 + 14 + 112 + 14).
    Use model_state_bytes.
    """
    # TODO
    raise NotImplementedError("Task 1c: rl_memory_gb")