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

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

class CLIPScoreNode:
    """Calculates true CLIP cosine similarity between prompt and image."""

    @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 (cached after first run)
            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]  # First from batch
            else:
                img = image
            
            # [H,W,C] float32 [0,1] → numpy uint8 → PIL RGB
            img_np = (img.cpu().numpy() * 255).astype(np.uint8)
            if img_np.shape[-1] == 4:  # RGBA → RGB
                img_np = img_np[:, :, :3]
            pil_image = Image.fromarray(img_np, mode='RGB')
            
            # GPU inference
            device = next(model.parameters()).device
            inputs = processor(text=[prompt], images=pil_image, return_tensors="pt", padding=True)
            inputs = {k: v.to(device) for k, v in inputs.items()}
            model.to(device)
            
            # Raw CLIP score: cosine similarity of normalized embeddings
            with torch.no_grad():
                outputs = model(**inputs)
                image_embeds = outputs.image_embeds
                text_embeds = outputs.text_embeds
                
                # L2 normalize + cosine similarity (standard CLIP score)
                image_embeds /= image_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-8
                text_embeds /= text_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-8
                score = (image_embeds @ text_embeds.T)[0, 0].item()
            
            # Cleanup
            model.cpu()
            del model, inputs
            
        except Exception as e:
            score = 0.0
            print(f"CLIP error: {str(e)}")
        
        score_str = f"CLIP Score: {score:.4f}"
        print(score_str)
        
        return (score, score_str)


NODE_CLASS_MAPPINGS = {
    "CLIPScoreNode": CLIPScoreNode,
}

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