import copy

import torch
from torch import nn


torch.manual_seed(7)

features = torch.linspace(-1.5, 1.5, steps=24, dtype=torch.float32).reshape(8, 3)
targets = features @ torch.tensor([[0.7], [-0.2], [0.4]]) + 0.15

model = nn.Sequential(nn.Linear(3, 1))
reference_model = copy.deepcopy(model)

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.05)

microbatch_size = 2
accumulation_steps = 2
effective_batch_size = microbatch_size * accumulation_steps

optimizer.zero_grad(set_to_none=True)

microbatch_count = 0
optimizer_steps = 0

for start in range(0, len(features), microbatch_size):
    batch_features = features[start : start + microbatch_size]
    batch_targets = targets[start : start + microbatch_size]

    raw_loss = loss_fn(model(batch_features), batch_targets)
    scaled_loss = raw_loss / accumulation_steps
    scaled_loss.backward()

    microbatch_count += 1
    grad_norm = torch.sqrt(
        sum(
            parameter.grad.detach().pow(2).sum()
            for parameter in model.parameters()
            if parameter.grad is not None
        )
    ).item()

    print(
        f"microbatch {microbatch_count}: "
        f"raw_loss={raw_loss.item():.4f}, "
        f"scaled_loss={scaled_loss.item():.4f}, "
        f"grad_norm={grad_norm:.4f}"
    )

    if microbatch_count % accumulation_steps == 0:
        optimizer.step()
        optimizer_steps += 1
        optimizer.zero_grad(set_to_none=True)
        gradients_clear = all(parameter.grad is None for parameter in model.parameters())
        print(
            f"optimizer step {optimizer_steps}: "
            f"after {accumulation_steps} backward passes, "
            f"gradients_clear={gradients_clear}"
        )

for start in range(0, len(features), effective_batch_size):
    batch_features = features[start : start + effective_batch_size]
    batch_targets = targets[start : start + effective_batch_size]

    reference_optimizer.zero_grad(set_to_none=True)
    reference_loss = loss_fn(reference_model(batch_features), batch_targets)
    reference_loss.backward()
    reference_optimizer.step()

max_delta = max(
    (parameter - reference_parameter).abs().max().item()
    for parameter, reference_parameter in zip(model.parameters(), reference_model.parameters())
)

print(f"optimizer steps taken: {optimizer_steps}")
print(f"full-batch max parameter difference: {max_delta:.8f}")