"""Interface to the [OpenCLIP](https://github.com/mlfoundations/open_clip) library."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/models/frozen_open_clip.ipynb.

# %% auto #0
__all__ = ['FrozenOpenCLIPEmbedderConfig', 'FrozenOpenCLIPEmbedder', 'CachedFrozenOpenCLIPEmbedderConfig',
           'CachedFrozenOpenCLIPEmbedder']

# %% ../../src/models/frozen_open_clip.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from .config_model import ConfigModel
from ..utils.async_fn import run_parallel_jobs
from ..utils.misc_utils import infer_torch_device
import open_clip

# %% ../../src/models/frozen_open_clip.ipynb #13bd5df0-fbf5-4f5e-bb82-6cab672c3542
@dataclass
class FrozenOpenCLIPEmbedderConfig:
    arch: str
    version: str
    #device: str
    max_length: int
    freeze: bool
    layer: str

# %% ../../src/models/frozen_open_clip.ipynb #dffe5b29-8362-46fe-9cdc-e83c996ca8b0
class FrozenOpenCLIPEmbedder(ConfigModel):
    """Loads and freezes the [OpenCLIP](https://github.com/mlfoundations/open_clip) transformer encoder for text prompts."""
    
    LAYERS = [
        # "pooled",
        "last",
        "penultimate"
    ]

    njobs = 1

    def __init__(self, arch="ViT-B-32", version="datacomp_xl_s13b_b90k", max_length=77, freeze=True, layer="penultimate", **kwargs):
        super().__init__()        
        
        assert layer in self.LAYERS     
        self.params_config = FrozenOpenCLIPEmbedderConfig(arch, version, max_length, freeze, layer)
        
        model, _, _ = open_clip.create_model_and_transforms(arch, device="cpu", pretrained=version)
        self.device = "cpu"
        
        del model.visual     
        self.model = model
        # self.to(device)
        
        self.tokenizer = open_clip.get_tokenizer(arch)
        assert torch.numel(self.tokenizer("test"))
        
        assert max_length <= 77   # max set by the clip 
        self.max_length = max_length
        
        if freeze: self.freeze()
        
        self.layer = layer
        if   self.layer == "last":         self.layer_idx = 0
        elif self.layer == "penultimate":  self.layer_idx = 1
        else: raise NotImplementedError()

        #create empty token, can also be, e.g., A nice picture
        self.empty_token = self.tokenize_and_push_to_device("")
        
    def freeze(self, freeze: bool = True):
        super().freeze(freeze=freeze)
                    
        for param in self.model.parameters(): 
            param.requires_grad = not freeze
    
    def to(self, device):
        self.model  = self.model.to(device)           
        self.device = device
        return self
        
    @torch.inference_mode()
    def tokenize_and_push_to_device(self, text, to_device=True):
        if self.njobs > 1:

            tokens_list = run_parallel_jobs(self.tokenizer, np.array_split(text, self.njobs), self.njobs)
            tokens      = torch.cat(tokens_list, dim=0)
            
        else:
            # tokens = open_clip.tokenize(text)
            tokens = self.tokenizer(text)
        
        if to_device:
            tokens = tokens.to(self.device)
        return tokens
    
    @torch.inference_mode()
    def forward(self, c, **kwargs):
        return self.encode_with_transformer(c)

    @torch.inference_mode()
    def encode_with_transformer(self, text):
        cast_dtype = self.model.transformer.get_cast_dtype()
        
        x = self.model.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]        
        x = x + self.model.positional_embedding[None, :x.shape[1]].to(cast_dtype)

        if not self.model.transformer.batch_first:
            x = x.permute(1, 0, 2)  # NLD -> LND
        
        x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)

        if not self.model.transformer.batch_first:
            x = x.permute(1, 0, 2)  # LND -> NLD
        
        x = self.model.ln_final(x)  # [batch_size, n_ctx, transformer.width]
        
        return x

    @torch.inference_mode()
    def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
        for i, r in enumerate(self.model.transformer.resblocks):
            if i == len(self.model.transformer.resblocks) - self.layer_idx:
                break
            #if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting():
                #x = checkpoint(r, x, attn_mask)
            #else:
            
            x = r(x, attn_mask=attn_mask)
            
        return x

    #--------------------------------------------------------------
    
    def get_config(self, save_path=None, without_metadata=False):
        return super().get_config(save_path=None, without_metadata=without_metadata)
    
    def store_model(self, config_path: str, save_path: str=None, without_metadata=False):        
        super().store_model(config_path, save_path=None, without_metadata=without_metadata)

    @staticmethod
    def from_config(config, device: torch.device, save_path: str=None):  
        config["save_path"] = None
        return ConfigModel.from_config(config, device, save_path=None)        

# %% ../../src/models/frozen_open_clip.ipynb #b06adb4f-56f7-4ca1-bd21-fabe060eba61
@dataclass
class CachedFrozenOpenCLIPEmbedderConfig(FrozenOpenCLIPEmbedderConfig):
    enable_cache_token_limit: bool

# %% ../../src/models/frozen_open_clip.ipynb #b57fe509-5765-422c-a1b5-5acc153e4c9b
class CachedFrozenOpenCLIPEmbedder(FrozenOpenCLIPEmbedder):
    """Adds caching support to `FrozenOpenCLIPEmbedder`."""

    def __init__(self, arch="ViT-B-32", version="datacomp_xl_s13b_b90k", max_length=77, freeze=True, layer="penultimate", enable_cache_token_limit: bool = True, **kwargs):
        super().__init__(arch=arch, version=version, max_length=max_length, freeze=freeze, layer=layer, **kwargs)  
        self.enable_cache_token_limit = enable_cache_token_limit

        self.params_config = CachedFrozenOpenCLIPEmbedderConfig(arch, version, max_length, freeze, layer, enable_cache_token_limit)
    
    def get_token_count(self, tokens, padding_token=0):
        # tokens .. [b, seq]
        collabsed_tokens = (tokens != padding_token).to(torch.int32)
        return torch.count_nonzero(collabsed_tokens, dim=-1)  # [b]
    
    @torch.inference_mode()
    def generate_cache(self, str_list: list=None, tokens=None, cached_empty_token_index=None, b_size=2048, y_on_cpu=False):       
        self.cached_empty_token_index = cached_empty_token_index       
        if exists(str_list): self.cached_tokens = self.tokenize_and_push_to_device(str_list)      
        elif exists(tokens): self.cached_tokens = tokens
        else: raise RuntimeError("please provide str_list or tokens")
        
        # note: we need to split the tokens in batches for forward pass, n gets large
        # cached_tokens     [n, 77]      ... int
        # cached_embeddings [n, 77, 512] ... float

        if self.enable_cache_token_limit:
            self.max_length = self.get_token_count(self.cached_tokens).max().item()
            self.params_config.max_length               = self.max_length
            self.params_config.enable_cache_token_limit = self.enable_cache_token_limit
            print(f"[INFO]: - `generate_cache` infered a TOKEN limit of {self.max_length}")

            #self.cached_tokens = self.cached_tokens[:, :self.max_length]
        
        n = self.cached_tokens.shape[0]
        
        n_chunks = int(np.ceil(n / b_size))
        
        in_device = self.cached_tokens.device
                
        last_ind = 0
        for i, cached_tokens in tqdm(enumerate(self.cached_tokens.chunk(n_chunks)), total=n_chunks):
            
            x = super().forward(cached_tokens.to(self.device))  # ... [b, seq, ch]
            
            if i == 0:
                mem = n * x.shape[1] * x.shape[2] * x.element_size() * 1e-9
                print(f"[INFO]: caching trying to allocate memory {(n, x.shape[1], x.shape[2])} on {'cpu' if y_on_cpu else self.device}, approx. {mem:.3f} GB")
                self.cached_embeddings = torch.zeros((n, x.shape[1], x.shape[2]), device="cpu" if y_on_cpu else self.device, dtype=x.dtype) # alloc huge memory !!
                
            self.cached_embeddings[last_ind:last_ind+x.shape[0]] = x.to(self.cached_embeddings.device)
            
            last_ind += x.shape[0]

        if self.enable_cache_token_limit:
            self.cached_embeddings = self.cached_embeddings[:, :self.max_length]
        
        if not y_on_cpu:
            self.cached_embeddings = self.cached_embeddings.to(in_device) 

    @torch.inference_mode()
    def look_up_cos_sim_cached_index(self, str_list: list=None, tokens=None):
        if exists(str_list): tokens = self.tokenize_and_push_to_device(str_list)      
        else: raise RuntimeError("please provide str_list or tokens")
                                         
        emb   = super().forward(tokens.to(self.device))
        c_emb = self.cached_embeddings
        #-----------------
        # do cos sim search
        
        emb   = emb.flatten(start_dim=1)   # [m, seq*ch]
        c_emb = c_emb.flatten(start_dim=1) # [n, seq*ch]

        norm_emb   =   emb / torch.linalg.vector_norm(  emb, dim=1, keepdim=True)
        norm_c_emb = c_emb / torch.linalg.vector_norm(c_emb, dim=1, keepdim=True) 
 
        sim     = torch.matmul(norm_c_emb, norm_emb.T) # matmul out is [n, m]
        max_idx = torch.argmax(sim, dim=0)             # reduce the c_emb dim, [m]
     
        return max_idx       
                            
    # @torch.inference_mode()
    def forward(self, c, **kwargs):  
        in_device = c.device
        
        if   c.dim() == 1: c_emb = self.cached_embeddings[c.to(self.cached_embeddings.device)].to(in_device)         #list of ints       
        elif c.dim() == 2: c_emb = super().forward(c.to(self.device))   #tokenized input      
        else: raise NotImplementedError("")

        if self.enable_cache_token_limit:
            c_emb = c_emb[:, :self.max_length]
        
        return c_emb
