from pathlib import Path
import shutil

import torch
from torch import nn
from torch.utils.tensorboard import SummaryWriter
from tensorboard.backend.event_processing import event_accumulator


torch.manual_seed(19)
log_dir = Path("runs/tensorboard-demo")
if log_dir.exists():
    shutil.rmtree(log_dir)

features = torch.tensor(
    [
        [-1.0, 0.2, 0.1],
        [-0.5, 0.5, 0.4],
        [0.0, -0.3, 0.9],
        [0.5, 0.7, -0.4],
        [1.0, -0.6, -0.1],
        [1.5, 0.4, 0.2],
    ],
    dtype=torch.float32,
)
targets = torch.tensor([0, 0, 1, 1, 1, 0])

model = nn.Sequential(
    nn.Linear(3, 6),
    nn.ReLU(),
    nn.Linear(6, 2),
)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.15)
writer = SummaryWriter(log_dir=log_dir)

for epoch in range(5):
    optimizer.zero_grad(set_to_none=True)
    logits = model(features)
    loss = loss_fn(logits, targets)
    loss.backward()
    optimizer.step()

    with torch.no_grad():
        predictions = model(features).argmax(dim=1)
        accuracy = (predictions == targets).float().mean().item()

    writer.add_scalar("Loss/train", loss.item(), epoch)
    writer.add_scalar("Accuracy/train", accuracy, epoch)
    print(f"epoch={epoch} loss={loss.item():.4f} accuracy={accuracy:.4f}")

writer.flush()
writer.close()

events = sorted(log_dir.glob("events.out.tfevents.*"))
reader = event_accumulator.EventAccumulator(str(log_dir))
reader.Reload()
scalar_tags = sorted(reader.Tags()["scalars"])
loss_points = reader.Scalars("Loss/train")
accuracy_points = reader.Scalars("Accuracy/train")

print(f"torch_version={torch.__version__}")
print(f"log_dir={log_dir}")
print(f"event_files={len(events)}")
for path in events:
    print(f"event_file={path.name}")
print(f"scalar_tags={','.join(scalar_tags)}")
print(f"loss_points={len(loss_points)}")
print(f"accuracy_points={len(accuracy_points)}")
print(f"last_loss={loss_points[-1].value:.4f}")
print(f"last_accuracy={accuracy_points[-1].value:.4f}")