from pathlib import Path
import shutil

import torch
from torch import nn
from torch.profiler import ProfilerActivity, profile, record_function, schedule, tensorboard_trace_handler


torch.manual_seed(23)

trace_dir = Path("profiler-traces")
if trace_dir.exists():
    shutil.rmtree(trace_dir)


class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(8, 16),
            nn.ReLU(),
            nn.Linear(16, 3),
        )

    def forward(self, inputs):
        return self.layers(inputs)


model = TinyClassifier()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
loss_fn = nn.CrossEntropyLoss()


def train_step(step):
    inputs = torch.randn(12, 8)
    targets = torch.tensor([0, 1, 2, 1, 0, 2, 2, 1, 0, 1, 2, 0])

    optimizer.zero_grad(set_to_none=True)

    with record_function("forward_and_loss"):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    with record_function("backward_and_step"):
        loss.backward()
        optimizer.step()

    print(f"step={step} loss={loss.item():.4f}")


with profile(
    activities=[ProfilerActivity.CPU],
    schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=tensorboard_trace_handler(str(trace_dir), worker_name="profiler-demo"),
    record_shapes=True,
    acc_events=True,
) as profiler:
    for step in range(5):
        train_step(step)
        profiler.step()

trace_files = sorted(trace_dir.glob("*.pt.trace.json"))

print(f"torch_version={torch.__version__}")
print(f"trace_dir={trace_dir}")
print(f"trace_file_count={len(trace_files)}")
for path in trace_files:
    print(f"trace_file={path.name}")

print(
    profiler.key_averages().table(
        sort_by="self_cpu_time_total",
        row_limit=6,
    )
)