import os
import numpy as np
from PIL import Image
from tqdm import tqdm
from sklearn.preprocessing import LabelEncoder
from transformers import CLIPProcessor, CLIPModel
import torch
from config import DEVICE

clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(DEVICE)
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
    
def extract_features_from_folder(folder_path):
    features, labels = [], []
    for label in os.listdir(folder_path):
        label_path = os.path.join(folder_path, label)
        if not os.path.isdir(label_path): continue

        for img_file in tqdm(os.listdir(label_path), desc=f"Processing {label}"):
            img_path = os.path.join(label_path, img_file)
            try:
                image = Image.open(img_path).convert("RGB")
                inputs = clip_processor(images=image, return_tensors="pt").to(DEVICE)
                with torch.no_grad():
                    img_feat = clip_model.get_image_features(**inputs)
                    # Handle newer transformers returning BaseModelOutputWithPooling
                    if hasattr(img_feat, 'pooler_output'):
                        img_feat = img_feat.pooler_output
                    img_feat = img_feat / img_feat.norm(p=2, dim=-1, keepdim=True)
                features.append(img_feat.cpu().numpy().squeeze())
                labels.append(label)
            except Exception as e:
                print(f"Error reading {img_path}: {e}")

    return np.array(features), labels

def prepare_dataset(train_dir, test_dir):
    X_train, y_train = extract_features_from_folder(train_dir)
    X_test, y_test = extract_features_from_folder(test_dir)

    encoder = LabelEncoder()
    y_train_enc = encoder.fit_transform(y_train)
    y_test_enc = encoder.transform(y_test)

    return X_train, y_train_enc, X_test, y_test_enc, encoder