import argparse

import torch
from torch import nn


def check_finite(name: str, tensor: torch.Tensor, *, epoch: int, batch: int) -> None:
    finite_mask = torch.isfinite(tensor)
    if bool(finite_mask.all()):
        return

    bad_values = tensor.detach()[~finite_mask].flatten()[:4].tolist()
    print(f"non_finite={name}")
    print(f"epoch={epoch} batch={batch} dtype={tensor.dtype} shape={tuple(tensor.shape)}")
    print(f"{name}_bad_values={bad_values}")
    print("hint=inspect the operation that created this tensor before changing the optimizer")
    raise SystemExit(1)


def check_gradients(model: nn.Module, *, epoch: int, batch: int) -> None:
    for name, parameter in model.named_parameters():
        if parameter.grad is None:
            continue
        if not bool(torch.isfinite(parameter.grad).all()):
            bad_values = parameter.grad.detach()[~torch.isfinite(parameter.grad)].flatten()[:4].tolist()
            print(f"non_finite_gradient={name}")
            print(f"epoch={epoch} batch={batch} bad_values={bad_values}")
            raise SystemExit(1)
    print("gradient_check=finite")


class TinyRegressor(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(2, 1)
        with torch.no_grad():
            self.linear.weight.copy_(torch.tensor([[-0.35, -0.20]]))
            self.linear.bias.fill_(-0.05)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.linear(inputs)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fixed", action="store_true")
    args = parser.parse_args()

    torch.manual_seed(7)
    torch.autograd.set_detect_anomaly(True, check_nan=True)

    features = torch.tensor(
        [
            [1.0, 0.5],
            [2.0, 1.0],
            [3.0, 1.5],
            [4.0, 2.0],
        ],
        dtype=torch.float32,
    )
    target = torch.tensor([[0.5], [1.0], [1.5], [2.0]], dtype=torch.float32)

    model = TinyRegressor()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

    epoch = 1
    batch = 0
    optimizer.zero_grad(set_to_none=True)
    prediction = model(features)
    check_finite("prediction", prediction, epoch=epoch, batch=batch)
    print(f"prediction_min={prediction.min().item():.6f}")
    print(f"prediction_max={prediction.max().item():.6f}")

    if args.fixed:
        loss = nn.functional.mse_loss(prediction, target)
    else:
        loss = torch.sqrt(prediction).mean()

    check_finite("loss", loss, epoch=epoch, batch=batch)
    loss.backward()
    check_gradients(model, epoch=epoch, batch=batch)
    optimizer.step()

    print(f"loss={loss.item():.6f}")
    print("optimizer_step=complete")


if __name__ == "__main__":
    main()