from .base_prompter import BasePrompter, tokenize_long_prompt
from ..models.model_manager import ModelManager, load_state_dict, search_for_embeddings
from ..models import SDTextEncoder
from transformers import CLIPTokenizer
import torch, os



class SDPrompter(BasePrompter):
    def __init__(self, tokenizer_path=None):
        if tokenizer_path is None:
            base_path = os.path.dirname(os.path.dirname(__file__))
            tokenizer_path = os.path.join(base_path, "tokenizer_configs/stable_diffusion/tokenizer")
        super().__init__()
        self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path)
        self.text_encoder: SDTextEncoder = None
        self.textual_inversion_dict = {}
        self.keyword_dict = {}


    def fetch_models(self, text_encoder: SDTextEncoder = None):
        self.text_encoder = text_encoder


    def add_textual_inversions_to_model(self, textual_inversion_dict, text_encoder):
        dtype = next(iter(text_encoder.parameters())).dtype
        state_dict = text_encoder.token_embedding.state_dict()
        token_embeddings = [state_dict["weight"]]
        for keyword in textual_inversion_dict:
            _, embeddings = textual_inversion_dict[keyword]
            token_embeddings.append(embeddings.to(dtype=dtype, device=token_embeddings[0].device))
        token_embeddings = torch.concat(token_embeddings, dim=0)
        state_dict["weight"] = token_embeddings
        text_encoder.token_embedding = torch.nn.Embedding(token_embeddings.shape[0], token_embeddings.shape[1])
        text_encoder.token_embedding = text_encoder.token_embedding.to(dtype=dtype, device=token_embeddings[0].device)
        text_encoder.token_embedding.load_state_dict(state_dict)


    def add_textual_inversions_to_tokenizer(self, textual_inversion_dict, tokenizer):
        additional_tokens = []
        for keyword in textual_inversion_dict:
            tokens, _ = textual_inversion_dict[keyword]
            additional_tokens += tokens
            self.keyword_dict[keyword] = " " + " ".join(tokens) + " "
        tokenizer.add_tokens(additional_tokens)


    def load_textual_inversions(self, model_paths):
        for model_path in model_paths:
            keyword = os.path.splitext(os.path.split(model_path)[-1])[0]
            state_dict = load_state_dict(model_path)

            # Search for embeddings
            for embeddings in search_for_embeddings(state_dict):
                if len(embeddings.shape) == 2 and embeddings.shape[1] == 768:
                    tokens = [f"{keyword}_{i}" for i in range(embeddings.shape[0])]
                    self.textual_inversion_dict[keyword] = (tokens, embeddings)

        self.add_textual_inversions_to_model(self.textual_inversion_dict, self.text_encoder)
        self.add_textual_inversions_to_tokenizer(self.textual_inversion_dict, self.tokenizer)


    def encode_prompt(self, prompt, clip_skip=1, device="cuda", positive=True):
        prompt = self.process_prompt(prompt, positive=positive)
        for keyword in self.keyword_dict:
            if keyword in prompt:
                print(f"Textual inversion {keyword} is enabled.")
                prompt = prompt.replace(keyword, self.keyword_dict[keyword])
        input_ids = tokenize_long_prompt(self.tokenizer, prompt).to(device)
        prompt_emb = self.text_encoder(input_ids, clip_skip=clip_skip)
        prompt_emb = prompt_emb.reshape((1, prompt_emb.shape[0]*prompt_emb.shape[1], -1))

        return prompt_emb