import torch
from transformers import AutoProcessor
from ..core import ModelConfig
from ..core.device.npu_compatible_device import get_device_type
from ..models.qwen_image_bench import QwenImageBenchModel
from .base import Metric
from transformers.utils import logging
logging.set_verbosity_error()


class QwenImageBenchMetric(Metric):
    def __init__(self, model: QwenImageBenchModel):
        super().__init__()
        self.model = model

    @classmethod
    def from_pretrained(
        cls,
        model_config: ModelConfig = ModelConfig(
            model_id="Qwen/Qwen-Image-Bench",
            origin_file_pattern="model-*.safetensors",
        ),
        processor_config: ModelConfig = ModelConfig(
            model_id="Qwen/Qwen-Image-Bench",
            origin_file_pattern="",
        ),
        torch_dtype: torch.dtype = None,
        device: torch.device = get_device_type(),
        max_new_tokens: int = 4096,
        resize_long_edge: int = 1024,
        processor_kwargs: dict = None,
        vram_limit: float = None,
    ):
        processor_kwargs = processor_kwargs or {}
        model_pool = cls.download_and_load_models(
            [model_config],
            torch_dtype=torch_dtype or torch.bfloat16,
            device=device,
            vram_limit=vram_limit,
        )
        model = model_pool.fetch_model("image_metrics_qwen_image_bench")
        if model is None:
            raise ValueError("Cannot find model: image_metrics_qwen_image_bench")
        if hasattr(model, "model"):
            model = model.model

        processor_config.download_if_necessary()
        processor = AutoProcessor.from_pretrained(processor_config.path, **processor_kwargs)
        model = QwenImageBenchModel(
            model=model,
            processor=processor,
            max_new_tokens=max_new_tokens,
            resize_long_edge=resize_long_edge,
        ).eval()
        return cls(model)

    @torch.no_grad()
    def evaluate(self, prompt: str | list[str] | None, images, dimensions=None):
        return self.model(prompt, images, dimensions=dimensions)

    @torch.no_grad()
    def score(self, prompt: str | list[str] | None, images, dimensions=None):
        outputs = self.evaluate(prompt, images, dimensions=dimensions)
        return [self.model._primary_score(output) for output in outputs]

    def compute(self, prompt: str | list[str] | None, images, dimensions=None):
        return self.score(prompt, images, dimensions=dimensions)

    def forward(self, prompt: str | list[str] | None, images, dimensions=None):
        return self.score(prompt, images, dimensions=dimensions)