import os
import cv2
import torch
import numpy as np
from torch.utils.data import DataLoader, TensorDataset
from torch import nn, optim
from BN import CatDogClassifier
import time

# Define image size
IMG_SIZE = 128

# Function to load images from a directory
def load_images_from_folder(folder, label):
    data = []
    for filename in os.listdir(folder):
        img_path = os.path.join(folder, filename)
        img = cv2.imread(img_path)  # Load image
        if img is not None:
            img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))  # Resize
            img = img / 255.0  # Normalize
            data.append((img, label))
    return data

# Load dataset
cat_data = load_images_from_folder("dataset/training_set/cats", label=0)  # Label 0 for cats
dog_data = load_images_from_folder("dataset/training_set/dogs", label=1)  # Label 1 for dogs

# Combine and shuffle data
dataset = cat_data + dog_data
np.random.shuffle(dataset)

# Convert to NumPy arrays
X = np.array([item[0] for item in dataset], dtype=np.float32)  # Image data
Y = np.array([item[1] for item in dataset], dtype=np.int64)  # Labels

# Reshape to PyTorch format (N, C, H, W)
X = np.transpose(X, (0, 3, 1, 2))  # Convert to (batch, channels, height, width)

# Convert to PyTorch tensors
X_tensor = torch.tensor(X).to("cuda")
Y_tensor = torch.tensor(Y).to("cuda")

# Create dataset and data loader
dataset = TensorDataset(X_tensor, Y_tensor)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# Initialize model, loss, and optimizer
model = CatDogClassifier()
model = model.to("cuda")
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
num_epochs = 25
start_time = time.time()
for epoch in range(num_epochs):
    total_loss = 0
    for images, labels in dataloader:
        optimizer.zero_grad()  # Reset gradients
        outputs = model(images)  # Forward pass
        loss = criterion(outputs, labels)  # Compute loss
        loss.backward()  # Backpropagation
        optimizer.step()  # Update weights
        total_loss += loss.item()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(dataloader):.4f}")

print(f"Time taken: {(time.time()-start_time):.2f} seconds")


torch.save(model, "super_model.pth")