import numpy as np
from pathlib import Path
from PIL import Image

import torch
import torchvision.transforms as T
import timm
from transformers import Blip2Processor, Blip2ForConditionalGeneration

class ImageProcessor:
    def __init__(self, weights_path: str | Path):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Модель извлечения визуальных признаков
        self.feature_extractor = timm.create_model('resnet50', pretrained=False, num_classes=8)
        
        if Path(weights_path).exists():
            self.feature_extractor.load_state_dict(torch.load(weights_path, map_location=self.device))
        else:
            print(f"Не удалось найти веса ResNet по пути: {weights_path}")
            
        # Удаление слоя классификации для вывода сырого вектора эмбеддингов
        self.feature_extractor.fc = torch.nn.Identity()
        self.feature_extractor.to(self.device).eval()

        # Трансформации для предварительной обработки изображений
        self.preprocess_image = T.Compose([
            T.Resize((224, 224)),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

        # Модуль семантического описания сцены
        print("Инициализация BLIP-2...")
        # Обход бага конфигурации Hugging Face (ручная сборка процессора)
        from transformers import BlipImageProcessor, AutoTokenizer
        img_proc = BlipImageProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
        tok = AutoTokenizer.from_pretrained("Salesforce/blip2-opt-2.7b", use_fast=False)
        self.blip_processor = Blip2Processor(image_processor=img_proc, tokenizer=tok)
        self.blip_model = Blip2ForConditionalGeneration.from_pretrained(
            "Salesforce/blip2-opt-2.7b", 
            torch_dtype=torch.float16 
        ).to(self.device)

    @torch.no_grad()
    def extract_embedding(self, image: Image.Image) -> np.ndarray:
        # Извлечение эмбеддингов из изображения
        rgb_image = image.convert('RGB')
        img_tensor = self.preprocess_image(rgb_image).unsqueeze(0).to(self.device)
        
        features = self.feature_extractor(img_tensor)
        features_np = features.cpu().numpy()
        
        return features_np.flatten()
        
    @torch.no_grad()
    def describe_scene(self, image: Image.Image) -> str:
        # Генерация текстового описания сцены
        rgb_image = image.convert('RGB')
        
        inputs = self.blip_processor(images=rgb_image, return_tensors="pt").to(self.device, torch.float16)
        generated_ids = self.blip_model.generate(**inputs, max_new_tokens=40)
        
        scene_description = self.blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        
        return scene_description.strip()