import torch
import numpy as np
import cv2
from transformers import AutoModelForImageClassification, AutoFeatureExtractor

class FitPredictor:
    def __init__(self):
        self.model_name = "facebook/dinov2-base"
        self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name)
        self.model = AutoModelForImageClassification.from_pretrained(self.model_name)
        self.model.eval()
        
        self.fit_types = ['tight', 'regular', 'loose', 'oversized']
        
        self.fit_recommendations = {
            'tight': ['Size up for comfort', 'Choose stretchy fabrics', 'Consider relaxed fit alternative'],
            'regular': ['True to size', 'Standard fit works for most body types'],
            'loose': ['Size down for fitted look', 'Great for layering', 'Can be styled with belt'],
            'oversized': ['Intentional oversized style', 'Style with fitted pieces', 'Works well for streetwear']
        }
        
        self.size_adjustments = {
            'tight': '+1',
            'regular': '0',
            'loose': '-1',
            'oversized': '-2'
        }
    
    def predict(self, image: np.ndarray) -> dict:
        inputs = self.feature_extractor(images=image, return_tensors="pt")
        
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=-1)
        
        probs = probabilities[0].tolist()
        
        fit_type = self.fit_types[0]
        confidence = 0.5
        for i, ft in enumerate(self.fit_types):
            if i < len(probs):
                if probs[i] > confidence:
                    confidence = probs[i]
                    fit_type = ft
        
        recommendations = self.fit_recommendations.get(fit_type, self.fit_recommendations['regular'])
        size_adjustment = self.size_adjustments.get(fit_type, '0')
        
        return {
            "fit": fit_type,
            "confidence": confidence,
            "recommendations": recommendations,
            "size_adjustment": size_adjustment,
            "body_compatibility": self._get_body_compatibility(fit_type)
        }
    
    def _get_body_compatibility(self, fit_type: str) -> dict:
        compatibility = {
            'tight': {'slim': 0.9, 'athletic': 0.8, 'curvy': 0.6, 'plus': 0.4},
            'regular': {'slim': 0.8, 'athletic': 0.9, 'curvy': 0.8, 'plus': 0.7},
            'loose': {'slim': 0.7, 'athletic': 0.8, 'curvy': 0.9, 'plus': 0.9},
            'oversized': {'slim': 0.9, 'athletic': 0.7, 'curvy': 0.6, 'plus': 0.5}
        }
        return compatibility.get(fit_type, compatibility['regular'])