import torch
from torch import nn


torch.manual_seed(7)


class SensorRegressor(nn.Module):
    def __init__(self, in_features=4, hidden_features=8, out_features=2):
        super().__init__()
        self.input = nn.Linear(in_features, hidden_features)
        self.activation = nn.ReLU()
        self.output = nn.Linear(hidden_features, out_features)

    def forward(self, features):
        hidden = self.activation(self.input(features))
        return self.output(hidden)


features = torch.tensor(
    [
        [0.2, 0.1, 0.7, 0.4],
        [0.9, 0.0, 0.5, 0.3],
        [0.4, 0.4, 0.2, 0.8],
    ],
    dtype=torch.float32,
)
targets = torch.tensor(
    [
        [0.5, 0.1],
        [0.8, 0.2],
        [0.3, 0.6],
    ],
    dtype=torch.float32,
)

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

predictions = model(features)
loss = loss_fn(predictions, targets)

optimizer.zero_grad(set_to_none=True)
loss.backward()

gradient_norms = [
    parameter.grad.detach().norm()
    for parameter in model.parameters()
    if parameter.grad is not None
]
gradient_norm = torch.linalg.vector_norm(torch.stack(gradient_norms), 2)

output_weight_before = model.output.weight.detach().clone()
optimizer.step()
weight_delta = (model.output.weight.detach() - output_weight_before).abs().max()

print(f"torch_version={torch.__version__}")
print("module=SensorRegressor")
print(f"parameter_tensors={len(list(model.parameters()))}")
print(f"trainable_parameters={sum(parameter.numel() for parameter in model.parameters())}")
print(f"output_shape={tuple(predictions.shape)}")
print(f"loss={loss.item():.6f}")
print(f"gradient_norm={gradient_norm.item():.6f}")
print(f"optimizer_step_changed_weight={bool(weight_delta > 0)}")