import os
import numpy as np
from tqdm import tqdm
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import Model
import os
import zipfile
import requests
from io import BytesIO
import gdown

def download_and_extract_dataset():
    if not os.path.exists("data"):
        print(" Downloading dataset with gdown...")
        file_id = "1jgcA9_rAEYw-JuGwepb9h1DVedWA2vUy"
        url = f"https://drive.google.com/file/d/1jgcA9_rAEYw-JuGwepb9h1DVedWA2vUy/view?usp=sharing"
        output = os.path.join("features", "sneakers_dataset.zip")
        os.makedirs("features", exist_ok=True)

        gdown.download(url, output, quiet=False)
        print(" Download complete. Extracting...")

        with zipfile.ZipFile(output, 'r') as zip_ref:
            zip_ref.extractall("data")
        print(" Dataset extracted to 'data/'")

download_and_extract_dataset()

# Directory paths
DATA_DIR = 'data'
FEATURE_DIR = 'tmp_features'

os.makedirs(FEATURE_DIR, exist_ok=True)
output_zip = os.path.join(FEATURES_DIR, "sneakers_dataset.zip")

# Load pre-trained ResNet50 model
base_model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
model = Model(inputs=base_model.input, outputs=base_model.output)

def preprocess_img(img_path):
    img = image.load_img(img_path, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    return preprocess_input(x)


def extract_features():
    features = []
    filenames = []

    # Recursively walk through subfolders
    for root, dirs, files in os.walk(DATA_DIR):
        for fname in files:
            if fname.lower().endswith(('.jpg', '.jpeg', '.png')):
                path = os.path.join(root, fname)
                try:
                    img_tensor = preprocess_img(path)
                    feature = model.predict(img_tensor, verbose=0)[0]
                    features.append(feature)
                    filenames.append(path)
                except Exception as e:
                    print(f"Error processing {path}: {e}")

    features = np.array(features)
    filenames = [f.replace("\\", "/") for f in filenames]
    filenames = np.array(filenames)
    
    np.save(os.path.join(FEATURE_DIR, 'features.npy'), features)
    np.save(os.path.join(FEATURE_DIR, 'filenames.npy'), filenames)

    print(f"Saved { len(features)} features from images in {DATA_DIR}/")

    print(f"✅ Example filename: {filenames[0]}")
if __name__ == "__main__":
    extract_features()