# Copyright (c) 2022 Machine Vision and Learning Group, LMU Munich
# copy from https://github.com/CompVis/latent-diffusion/blob/main/ldm/lr_scheduler.py under the MIT license
#   LICENSE is in incl_licenses directory.
import torch
import importlib
import torch.nn as nn

from inspect import isfunction

from transformers import (
    AutoTokenizer,
    T5EncoderModel,
)


def exists(x):
    return x is not None


def default(val, d):
    if exists(val):
        return val
    return d() if isfunction(d) else d


def instantiate_from_config(config, reload=False):
    if not "target" in config:
        if config == "__is_first_stage__":
            return None
        elif config == "__is_unconditional__":
            return None
        raise KeyError("Expected key `target` to instantiate.")
    return get_obj_from_str(config["target"], reload=reload)(
        **config.get("params", dict())
    )


def get_obj_from_str(string, reload=False):
    module, cls = string.rsplit(".", 1)
    if reload:
        module_imp = importlib.import_module(module)
        importlib.reload(module_imp)
    return getattr(importlib.import_module(module, package=None), cls)


class TextEncoder(nn.Module):

    def __init__(self, ckpt_path=None):
        super().__init__()

        self.tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
        self.model = T5EncoderModel.from_pretrained(ckpt_path)
        self.model.eval()
        self.device = "cuda"

    def get_uncond(self, batchsize):

        hidden_state, _ = self([""])
        hidden_state = torch.cat([hidden_state] * batchsize).float()
        attention_mask = (
            torch.ones((batchsize, hidden_state.size(1))).to(self.device).float()
        )
        return [hidden_state, attention_mask]

    def forward(self, batch):
        return self.get_cond(batch)

    def get_cond(self, prompt):
        device = self.model.device
        token = self.tokenizer(
            prompt,
            max_length=128,
            padding=True,
            truncation=True,
            return_tensors="pt",
        )
        input_ids, attention_mask = token.input_ids.to(device), token.attention_mask.to(
            device
        )

        with torch.no_grad():
            hidden_states = self.model(
                input_ids=input_ids, attention_mask=attention_mask
            )[0]

        return [hidden_states, attention_mask.float()]