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

IMG_SIZE = 128

if torch.cuda.is_available():
    DEVICE = torch.device("cuda")
elif torch.mps.is_available():
    DEVICE = torch.device("mps")
else:
    DEVICE = torch.device("cpu")

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)
        if img is not None:
            img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))
            img = img / 255.0
            data.append((img, label))
    return data

def train_model(model_name, spinner):

    spinner.write(f'Using the following device : {DEVICE}')

    # Loading the dataset
    cat_data = load_images_from_folder("dataset/training_set/cats", label=0)
    dog_data = load_images_from_folder("dataset/training_set/dogs", label=1)

    dataset = cat_data + dog_data
    np.random.shuffle(dataset)

    X = np.array([item[0] for item in dataset], dtype=np.float32)
    Y = np.array([item[1] for item in dataset], dtype=np.int64)

    X = np.transpose(X, (0, 3, 1, 2))

    X_tensor = torch.tensor(X).to(DEVICE)
    Y_tensor = torch.tensor(Y).to(DEVICE)

    dataset = TensorDataset(X_tensor, Y_tensor)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

    model = CatDogClassifier(img_size=IMG_SIZE)
    model = model.to(DEVICE)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    num_epochs = 20
    start_time = time.time()

    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        for images, labels in dataloader:
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

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

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

    torch.save(model, f"models/{model_name}.pth")