import torch
from torch import nn


class WeightedHuberLoss(nn.Module):
    def __init__(self, delta=0.5, under_prediction_weight=2.0):
        super().__init__()
        self.delta = delta
        self.under_prediction_weight = under_prediction_weight

    def forward(self, prediction, target):
        error = prediction - target
        abs_error = error.abs()
        quadratic = torch.minimum(abs_error, torch.full_like(abs_error, self.delta))
        linear = abs_error - quadratic
        huber = 0.5 * quadratic.pow(2) + self.delta * linear
        weights = torch.where(
            error < 0,
            torch.full_like(error, self.under_prediction_weight),
            torch.ones_like(error),
        )
        return (weights * huber).mean()


torch.manual_seed(19)

features = torch.tensor(
    [
        [0.0, 0.5],
        [1.0, -0.5],
        [2.0, 1.0],
        [-1.0, 1.5],
    ],
    dtype=torch.float32,
)
targets = torch.tensor([[0.4], [1.0], [2.4], [-0.2]], dtype=torch.float32)

model = nn.Sequential(
    nn.Linear(2, 4),
    nn.Tanh(),
    nn.Linear(4, 1),
)
loss_fn = WeightedHuberLoss(delta=0.5, under_prediction_weight=2.5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

optimizer.zero_grad(set_to_none=True)
prediction = model(features)
loss = loss_fn(prediction, targets)
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)
first_layer_grad_present = model[0].weight.grad is not None

first_layer_before = model[0].weight.detach().clone()
optimizer.step()
weight_delta = (model[0].weight.detach() - first_layer_before).abs().max()

print(f"torch_version={torch.__version__}")
print(f"custom_loss={loss.item():.6f}")
print(f"gradient_norm={gradient_norm.item():.6f}")
print(f"first_layer_grad_present={first_layer_grad_present}")
print(f"optimizer_step_changed_weight={weight_delta.item() > 0}")