import torch
import torch.nn as nn
from transformers import CLIPTextModel, CLIPTokenizer


class FrozenCLIPEmbedder(nn.Module):
    def __init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77, freeze=True, layer="last", layer_idx=None):
        super().__init__()
        self.tokenizer = CLIPTokenizer.from_pretrained(version)
        self.transformer = CLIPTextModel.from_pretrained(version)
        self.device = device
        self.max_length = max_length
        self.layer = layer
        self.layer_idx = layer_idx
        if freeze:
            self.transformer = self.transformer.eval()
            for p in self.parameters():
                p.requires_grad = False

    def forward(self, text):
        enc = self.tokenizer(
            text, truncation=True, max_length=self.max_length,
            return_length=True, return_overflowing_tokens=False,
            padding="max_length", return_tensors="pt"
        )
        tokens = enc["input_ids"].to(next(self.transformer.parameters()).device)
        out = self.transformer(input_ids=tokens, output_hidden_states=self.layer == "hidden")
        if self.layer == "last":
            return out.last_hidden_state
        if self.layer == "pooled":
            return out.pooler_output[:, None, :]
        return out.hidden_states[self.layer_idx]

    def encode(self, text):
        return self(text)