"""
ComfyUI Custom Node: CLIP Score Calculator
"""

import torch
import numpy as np
from PIL import Image
from transformers import CLIPProcessor, CLIPModel


class CLIPScoreNode:
    """Calculates CLIP cosine similarity with simple prompt chunking."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model_name": (
                    [
                        "openai/clip-vit-base-patch16",
                        "openai/clip-vit-base-patch32",
                        "openai/clip-vit-large-patch14",
                    ],
                    {"default": "openai/clip-vit-base-patch16"},
                ),
            },
        }

    RETURN_TYPES = ("FLOAT", "STRING")
    RETURN_NAMES = ("clip_score", "score_text")
    FUNCTION = "calculate_clip_score"
    CATEGORY = "metrics"
    OUTPUT_NODE = True

    def calculate_clip_score(self, image, prompt, model_name):
        try:
            # Load CLIP
            model = CLIPModel.from_pretrained(model_name)
            processor = CLIPProcessor.from_pretrained(model_name)

            # Handle ComfyUI IMAGE: [B,H,W,C] or [H,W,C]
            if len(image.shape) == 4:
                img = image[0]
            else:
                img = image

            # Convert to PIL
            img_np = (img.cpu().numpy() * 255).astype(np.uint8)
            if img_np.shape[-1] == 4:
                img_np = img_np[:, :, :3]
            pil_image = Image.fromarray(img_np, mode="RGB")

            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            model.to(device)

            chunks = [prompt[i:i + 200] for i in range(0, len(prompt), 200)]
            scores = []

            with torch.no_grad():
                for chunk in chunks:
                    inputs = processor(
                        text=[chunk],
                        images=pil_image,
                        return_tensors="pt",
                        padding=True,
                        truncation=True,
                        max_length=77,
                    )

                    inputs = {k: v.to(device) for k, v in inputs.items()}

                    outputs = model(**inputs)

                    image_embeds = outputs.image_embeds
                    text_embeds = outputs.text_embeds

                    # Normalize
                    image_embeds = image_embeds / (
                        image_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-8
                    )
                    text_embeds = text_embeds / (
                        text_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-8
                    )

                    score = (image_embeds @ text_embeds.T)[0, 0].item()
                    scores.append(score)

            final_score = sum(scores) / len(scores) if scores else 0.0

            # Cleanup
            model.cpu()
            del model

        except Exception as e:
            final_score = 0.0
            print(f"CLIP error: {str(e)}")

        score_str = f"CLIP Score: {final_score:.4f}"
        print(score_str)

        return (final_score, score_str)


NODE_CLASS_MAPPINGS = {
    "CLIPScoreNode": CLIPScoreNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "CLIPScoreNode": "CLIP Score Calculator",
}