import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader

# Ustawienia parametrów treningu
img_width, img_height = 224, 224  # Wymiary obrazu wymagane przez model ResNet
batch_size = 32  # Liczba obrazów przetwarzanych na raz podczas treningu
epochs = 10  # Liczba epok treningu
learning_rate = 0.001  # Wskaźnik uczenia się dla optymalizatora
model_path = './result/animal_classifier_resnet.pth'  # Ścieżka do zapisu wytrenowanego modelu

# Sprawdzenie, czy jest dostępny GPU i przypisanie urządzenia do zmiennej `device`
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Wyłączenie CuDNN, chyba że twoja karta wspiera bibliotekę CuDNN (NVIDIA CUDA Deep Neural Network library)
# torch.backends.cudnn.enabled = False

# Transformacje danych wejściowych (zmiana rozmiaru, normalizacja)
transform = transforms.Compose([
    transforms.Resize((img_width, img_height)),  # Zmiana rozmiaru obrazu
    transforms.ToTensor(),  # Konwersja obrazu do tensoru
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # Normalizacja obrazu
])

# Przygotowanie danych treningowych z katalogu `raw-img`
data_dir = 'raw-img'  # Ścieżka do katalogu z obrazami treningowymi
train_dataset = datasets.ImageFolder(data_dir, transform=transform)  # Wczytanie obrazów i zastosowanie transformacji
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)  # Tworzenie loadera danych

# Użycie pretrenowanego modelu ResNet18
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
num_ftrs = model.fc.in_features  # Liczba wejściowych cech ostatniej warstwy
model.fc = nn.Linear(num_ftrs, len(train_dataset.classes))  # Zastąpienie ostatniej warstwy dopasowanej do liczby klas w danych

# Przeniesienie modelu na GPU, jeśli jest dostępny
model = model.to(device)

# Definicja funkcji kosztu (CrossEntropyLoss) i optymalizatora (Adam)
criterion = nn.CrossEntropyLoss()  # Funkcja kosztu
optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # Optymalizator

# Trening modelu
for epoch in range(epochs):  # Pętla przez epoki
    model.train()  # Ustawienie modelu w tryb treningowy
    running_loss = 0.0  # Zmienna do śledzenia straty
    for inputs, labels in train_loader:  # Pętla przez batch'e danych
        # Przeniesienie danych na GPU, jeśli jest dostępny
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()  # Zerowanie gradientów
        outputs = model(inputs)  # Przekazanie danych przez model
        loss = criterion(outputs, labels)  # Obliczenie straty
        loss.backward()  # Propagacja wsteczna
        optimizer.step()  # Aktualizacja wag modelu
        running_loss += loss.item()  # Akumulacja straty

    print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader)}")  # Wyświetlenie średniej straty na epokę

# Zapisywanie wytrenowanego modelu do pliku
torch.save(model.state_dict(), model_path)
print(f"Model saved to {model_path}")