import torch
from torch import nn


class TinyXorNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(2, 8),
            nn.ReLU(),
            nn.Linear(8, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)


def main() -> None:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"torch={torch.__version__} cuda={torch.version.cuda} device={device}")
    if device == "cuda":
        print(f"gpu={torch.cuda.get_device_name(0)}")

    x = torch.tensor(
        [[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]],
        device=device,
    )
    y = torch.tensor([[0.0], [1.0], [1.0], [0.0]], device=device)

    model = TinyXorNet().to(device)
    loss_fn = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

    for step in range(1, 501):
        optimizer.zero_grad()
        logits = model(x)
        loss = loss_fn(logits, y)
        loss.backward()
        optimizer.step()

        if step % 100 == 0:
            print(f"step={step} loss={loss.item():.4f}")

    with torch.no_grad():
        probabilities = torch.sigmoid(model(x))
        predictions = (probabilities >= 0.5).int()

    print("predictions:")
    rows = zip(
        x.cpu().tolist(),
        y.cpu().int().tolist(),
        predictions.cpu().tolist(),
    )
    for inputs, expected, actual in rows:
        print(f"  {inputs} -> expected={expected[0]} predicted={actual[0]}")


if __name__ == "__main__":
    main()