import time
import torch
import torch.nn as nn
import torch.nn.functional as F


def train_step_with_accumulation(
    model, optimizer, micro_batches, accumulation_steps
):
    """Simulate a large effective batch via gradient accumulation."""
    optimizer.zero_grad()
    total_loss = 0.0
    for xb, yb in micro_batches:
        logits = model(xb)
        loss = F.cross_entropy(logits, yb) / accumulation_steps
        loss.backward()  # gradients accumulate in .grad by default
        total_loss += loss.item()
    optimizer.step()
    return total_loss


def profile_forward_backward(model, x, y, n_repeats=20):
    # warm-up (first call often includes one-time setup costs)
    logits = model(x)
    loss = F.cross_entropy(logits, y)
    loss.backward()
    model.zero_grad()

    start = time.perf_counter()
    for _ in range(n_repeats):
        logits = model(x)
        loss = F.cross_entropy(logits, y)
        loss.backward()
        model.zero_grad()
    elapsed = time.perf_counter() - start
    return elapsed / n_repeats


class CheckpointedMLP(nn.Module):
    """Illustrates activation checkpointing on a deep stack of
    layers using torch.utils.checkpoint."""

    def __init__(self, width, depth):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(width, width) for _ in range(depth)
        ])

    def forward(self, x, use_checkpoint=False):
        from torch.utils.checkpoint import checkpoint

        for layer in self.layers:
            if use_checkpoint:
                x = checkpoint(
                    lambda inp, l=layer: F.relu(l(inp)),
                    x, use_reentrant=False,
                )
            else:
                x = F.relu(layer(x))
        return x
