import torch
import torch.nn.functional as F
from transformers import ViTForImageClassification, ViTImageProcessor
import numpy as np
from PIL import Image
import json

class ClothingClassifier:
    def __init__(self):
        self.model_name = "google/vit-base-patch16-224"
        self.processor = ViTImageProcessor.from_pretrained(self.model_name)
        self.model = ViTForImageClassification.from_pretrained(self.model_name)
        self.model.eval()
        
        self.clothing_mapping = {
            'shirt': ['shirt', 't-shirt', 'blouse', 'top', 'polo'],
            'pants': ['pants', 'trousers', 'jeans', 'slacks', 'chinos'],
            'dress': ['dress', 'gown', 'frock', 'sundress'],
            'jacket': ['jacket', 'coat', 'blazer', 'hoodie', 'cardigan'],
            'shoes': ['shoe', 'sneakers', 'boots', 'sandals', 'loafers'],
            'hat': ['hat', 'cap', 'beanie', 'beret', 'fedora'],
            'skirt': ['skirt', 'mini', 'midi', 'maxi'],
            'shorts': ['shorts', 'bermuda', 'cutoffs'],
            'sweater': ['sweater', 'jumper', 'pullover', 'knitwear'],
            'accessory': ['bag', 'belt', 'scarf', 'tie', 'gloves']
        }
        
        self.all_labels = []
        for category, keywords in self.clothing_mapping.items():
            self.all_labels.extend(keywords)
    
    def predict(self, image: np.ndarray) -> dict:
        inputs = self.processor(images=image, return_tensors="pt")
        
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = F.softmax(logits, dim=-1)
        
        probs = probabilities[0].tolist()
        top_k = 10
        top_indices = torch.topk(probabilities[0], k=top_k).indices.tolist()
        
        predictions = []
        confidence_scores = []
        
        for idx in top_indices:
            label = self.model.config.id2label[idx].lower()
            confidence = probs[idx]
            predictions.append(label)
            confidence_scores.append(confidence)
        
        detected_type = self._map_to_clothing_type(predictions)
        confidence = max(confidence_scores) if confidence_scores else 0.0
        
        return {
            "type": detected_type,
            "confidence": confidence,
            "predictions": predictions[:5],
            "confidence_scores": confidence_scores[:5]
        }
    
    def _map_to_clothing_type(self, predictions: list) -> str:
        for pred in predictions:
            for category, keywords in self.clothing_mapping.items():
                if any(keyword in pred.lower() for keyword in keywords):
                    return category
        return "clothing"
    
    def batch_predict(self, images: list) -> list:
        results = []
        for img in images:
            results.append(self.predict(img))
        return results