import torch
import torch.nn.functional as F
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import numpy as np
import cv2

class TextureAnalyzer:
    def __init__(self):
        self.model_name = "microsoft/resnet-50"
        self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name)
        self.model = AutoModelForImageClassification.from_ppretrained(self.model_name)
        self.model.eval()
        
        self.texture_types = [
            'cotton', 'denim', 'leather', 'silk', 'wool', 'linen',
            'polyester', 'velvet', 'knit', 'mesh', 'lace', 'satin',
            'chiffon', 'corduroy', 'fleece', 'tweed', 'cashmere'
        ]
        
        self.texture_properties = {
            'cotton': {'breathable': True, 'stretchy': False, 'warm': False, 'durable': True},
            'denim': {'breathable': False, 'stretchy': False, 'warm': True, 'durable': True},
            'leather': {'breathable': False, 'stretchy': False, 'warm': True, 'durable': True},
            'silk': {'breathable': True, 'stretchy': False, 'warm': False, 'durable': False},
            'wool': {'breathable': True, 'stretchy': True, 'warm': True, 'durable': True},
            'linen': {'breathable': True, 'stretchy': False, 'warm': False, 'durable': True},
            'polyester': {'breathable': False, 'stretchy': True, 'warm': True, 'durable': True},
            'velvet': {'breathable': False, 'stretchy': False, 'warm': True, 'durable': False}
        }
        
        self.fabric_recommendations = {
            'cotton': ['Great for everyday wear', 'Breathable for summer'],
            'leather': ['Best for outerwear', 'Requires special care'],
            'silk': ['Luxurious feel', 'Dry clean recommended'],
            'wool': ['Perfect for winter', 'Keep away from moths']
        }
    
    def analyze(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 = F.softmax(logits, dim=-1)
        
        probs = probabilities[0].tolist()
        top_indices = torch.topk(probabilities[0], k=5).indices.tolist()
        
        texture = self.texture_types[top_indices[0]] if top_indices else 'cotton'
        confidence = probs[top_indices[0]] if top_indices else 0.5
        
        properties = self.texture_properties.get(texture, {'breathable': False, 'stretchy': False, 'warm': False, 'durable': False})
        recommendations = self.fabric_recommendations.get(texture, ['Versatile fabric for various occasions'])
        
        return {
            "texture": texture,
            "confidence": confidence,
            "properties": properties,
            "recommendations": recommendations
        }
    
    def extract_texture_features(self, image: np.ndarray) -> dict:
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        contrast = gray.std()
        smoothness = 1 - (gray.var() / (gray.max() + 1e-6))
        
        return {
            "contrast": float(contrast),
            "smoothness": float(smoothness),
            "roughness": float(1 - smoothness)
        }