import cv2
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
import torchvision.transforms as transforms

class ContentAnalyzer:
    def __init__(self):
        # Use lightweight model (MobileNetV2) - works on CPU
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = self._load_model()
        
    def _load_model(self):
        # Load pre-trained MobileNetV2 for feature extraction
        model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)
        # Remove classification head to get features
        model.classifier = nn.Identity()
        model.eval()
        return model.to(self.device)
    
    def analyze(self, image_path):
        """Returns image type and optimal compression strategy"""
        # Load image
        img = cv2.imread(str(image_path))
        if img is None:
            return "photo", {"method": "avif", "quality": 75}
        
        h, w = img.shape[:2]
        
        # 1. Edge detection (for text/screenshots)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150)
        edge_ratio = np.sum(edges > 0) / edges.size
        
        # 2. Color analysis
        unique_colors = len(np.unique(img.reshape(-1, img.shape[2]), axis=0))
        
        # 3. Texture analysis (variance of Laplacian)
        laplacian_var = cv2.Laplacian(gray, cv2.CV_64F).var()
        
        # 4. Check for transparency
        has_transparency = False
        try:
            pil_img = Image.open(image_path)
            has_transparency = pil_img.mode in ('RGBA', 'LA', 'P') and 'transparency' in pil_img.info
        except:
            pass
        
        # Decision logic
        if has_transparency:
            return "graphic_with_transparency", {
                "method": "png_optimized",
                "colors": 256,
                "lossless": True
            }
        elif edge_ratio > 0.15 and unique_colors < 5000:
            return "screenshot_or_text", {
                "method": "webp_lossless",
                "quality": 90,
                "preserve_text": True
            }
        elif unique_colors < 1000:
            return "graphic", {
                "method": "png_quantized",
                "colors": min(256, unique_colors),
                "dither": 0.5
            }
        elif laplacian_var < 100:
            return "smooth_gradient", {
                "method": "avif",
                "quality": 80,
                "avoid_banding": True
            }
        else:
            return "photo", {
                "method": "avif",
                "quality": 75,
                "chroma_subsampling": "4:2:0"
            }