import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


torch.manual_seed(23)

features = torch.tensor(
    [
        [-1.0, 0.0, 0.5],
        [-0.5, 0.25, 1.0],
        [0.0, -0.5, 0.25],
        [0.5, 0.75, -0.25],
        [1.0, -0.25, -0.5],
        [1.5, 0.5, 0.0],
        [2.0, -0.75, 0.75],
        [2.5, 1.0, -1.0],
    ],
    dtype=torch.float32,
)
targets = features @ torch.tensor([[0.8], [-0.4], [0.3]]) + 0.2

train_data = TensorDataset(features, targets)
train_loader = DataLoader(train_data, batch_size=4, shuffle=False)

model = nn.Sequential(nn.Linear(3, 1))
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.08)

first_weight = model[0].weight
weight_before = first_weight.detach().clone()

model.train()
epoch_losses = []

for epoch in range(1, 7):
    loss_total = 0.0
    sample_count = 0

    for batch_features, batch_targets in train_loader:
        optimizer.zero_grad(set_to_none=True)
        predictions = model(batch_features)
        loss = loss_fn(predictions, batch_targets)
        loss.backward()
        optimizer.step()

        loss_total += loss.item() * batch_features.size(0)
        sample_count += batch_features.size(0)

    average_loss = loss_total / sample_count
    epoch_losses.append(average_loss)
    print(f"epoch={epoch} loss={average_loss:.6f}")

with torch.inference_mode():
    sample_prediction = model(features[:2])

weight_changed = not torch.equal(weight_before, first_weight.detach())
loss_decreased = epoch_losses[-1] < epoch_losses[0]

print(f"weight_changed={weight_changed}")
print(f"loss_decreased={loss_decreased}")
print(f"prediction_shape={tuple(sample_prediction.shape)}")