import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import random  # For mock data generation

# Load Pretrained ResNet-50 Model
device = "cuda" if torch.cuda.is_available() else "cpu"
model = models.resnet50(pretrained=True).to(device)
model.eval()

# Remove the last fully connected layer
model = torch.nn.Sequential(*(list(model.children())[:-1]))

# Define Preprocessing Transform
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def extract_features(image_path):
    """Extracts ResNet feature vector from an image."""
    image = Image.open(image_path).convert("RGB")
    image = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        features = model(image).squeeze()
    
    return features.cpu().numpy().flatten().tolist()  # Convert tensor to list

def generate_image_description(image_path):
    """
    Mock function to generate clothing description based on image
    In a production environment, this would use a more sophisticated 
    image recognition model or an API like Google Vision
    """
    # Mock description generator - would be replaced with actual ML model
    neck_types = ["round-neck", "v-neck", "crew-neck", "turtle-neck", "boat-neck"]
    sleeve_types = ["half-sleeves", "full-sleeves", "sleeveless", "3/4 sleeves"]
    clothing_types = ["t-shirt", "shirt", "blouse", "sweater", "hoodie", "jacket"]
    
    neck = random.choice(neck_types)
    sleeve = random.choice(sleeve_types)
    clothing = random.choice(clothing_types)
    
    description = f"{neck} {sleeve} {clothing}"
    return description

def analyze_product_image(product_image_path, model_image_path=None):
    """
    Extract features from product image and generate description
    If model image is provided, extract features from that as well
    """
    product_features = extract_features(product_image_path)
    description = generate_image_description(product_image_path)
    
    # If we have both product and model images, we'll use product image features for search
    # In a more advanced system, you might want to combine features from both images
    
    return {
        "features": product_features,
        "auto_description": description
    }