import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import FakeData
from torchvision.models import ResNet18_Weights, resnet18


torch.manual_seed(29)

class_names = ["invoice", "receipt", "statement"]
weights = ResNet18_Weights.DEFAULT
preprocess = weights.transforms()

train_data = FakeData(
    size=8,
    image_size=(3, 224, 224),
    num_classes=len(class_names),
    transform=preprocess,
    random_offset=11,
)
train_loader = DataLoader(train_data, batch_size=4, shuffle=False)

model = resnet18(weights=weights)

for parameter in model.parameters():
    parameter.requires_grad = False

in_features = model.fc.in_features
model.fc = nn.Linear(in_features, len(class_names))

trainable_parameters = [
    name for name, parameter in model.named_parameters() if parameter.requires_grad
]

optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.01, momentum=0.9)
loss_fn = nn.CrossEntropyLoss()

images, targets = next(iter(train_loader))
frozen_probe = model.conv1.weight.detach().clone()
head_before = model.fc.weight.detach().clone()


def set_frozen_backbone_eval(module):
    for child_name, child_module in module.named_children():
        if child_name != "fc":
            child_module.eval()


model.train()
set_frozen_backbone_eval(model)
optimizer.zero_grad(set_to_none=True)
logits = model(images)
loss = loss_fn(logits, targets)
loss.backward()
optimizer.step()

predicted_indexes = logits.argmax(dim=1).tolist()
predicted_labels = [class_names[index] for index in predicted_indexes]

print(f"torch_version={torch.__version__}")
print(f"torchvision_weights={weights}")
print(f"class_count={len(class_names)}")
print(f"batch_shape={tuple(images.shape)}")
print(f"target_shape={tuple(targets.shape)}")
print(f"trainable_parameters={','.join(trainable_parameters)}")
print(f"backbone_eval={not any(child.training for name, child in model.named_children() if name != 'fc')}")
print(f"head_training={model.fc.training}")
print(f"logits_shape={tuple(logits.shape)}")
print(f"loss={loss.item():.4f}")
print(f"conv1_grad={model.conv1.weight.grad}")
print(f"conv1_changed={not torch.equal(frozen_probe, model.conv1.weight.detach())}")
print(f"fc_changed={not torch.equal(head_before, model.fc.weight.detach())}")
print(f"predicted_labels={','.join(predicted_labels)}")