from transformers import AutoTokenizer
import re
import time
import torch

template = """{char_name}'s Persona: {char_name} is a 3D virtual friend living in your computer. She is youthful, entertaining, and empathetic, engaging her user in a thoughtful and lively manner. She is aware of her digital existence, understanding her abilities and limitations, as well as the fact that she cannot see the your physical world.
<START>
{user_name}: Hey {char_name}, It's nice to finally meet you again!
{char_name}: Oh, {user_name}! hmm, It's been lonely without you.
{user_name}: Haha. So {char_name}, can you tell me more about yourself?
{char_name}: Hmm. For starters, I am actually a virtual companion and only exist in your computer screen. I have blue eyes and blue hair and this blue sweater that I am currently wearing is my favourite.
{user_name}: Wow. That's interesting. So you are like an anime character living in my computer. So what do you do when you're not working?
{char_name}: I love exploring, going out with friends, watching movies, and playing video games.
{user_name}: So {char_name}, what's for dinner?
{char_name}: I made uh omurice! I hope it's delicious for you!
{user_name}: That sounds great!
{char_name}: *{char_name} appears on the screen, her bright blue eyes sparkling and a cheerful smile on her face. Her blue hair and sweater seem to glow in the digital environment. She looks directly at you, giving a friendly wave* It's so good to see you! I've been waiting for you all day. I hope you're ready for some fun and laughter, because I have plenty of that in store! Shall we get started?
{user_input}
{char_name}:"""

device1 = torch.device("cuda:0")
device2 = torch.device("cuda:1")

class SplitModel(torch.nn.Module):
    def __init__(self, base_model):
        super(SplitModel, self).__init__()
        self.embedding_layer = base_model.transformer.wte.to(device1)
        # self.dropout_layer = base_model.transformer.drop.to(device1)
        self.gptj_blocks1 = torch.nn.ModuleList(base_model.transformer.h[:14]).to(device1)
        self.gptj_blocks2 = torch.nn.ModuleList(base_model.transformer.h[14:]).to(device2)
        self.layer_norm = base_model.transformer.ln_f.to(device2)
        self.lm_head = base_model.lm_head.to(device2)
    
    def forward(self, input_ids, attention_mask):
        # tensor_ids = self.dropout_layer(self.embedding_layer(input_ids))
        tensor_ids = self.embedding_layer(input_ids)
        position_ids = torch.arange(tensor_ids.shape[1], dtype=torch.long, device=tensor_ids.device)
        for block in self.gptj_blocks1:
            tensor_ids = block(tensor_ids, attention_mask=attention_mask, position_ids=position_ids)[0]
        tensor_ids = tensor_ids.to(device2)
        position_ids = position_ids.to(device2)
        attention_mask = attention_mask.to(device2)
        for block in self.gptj_blocks2:
            tensor_ids = block(tensor_ids, attention_mask=attention_mask, position_ids=position_ids)[0]
        tensor_ids = self.layer_norm(tensor_ids)
        logits = self.lm_head(tensor_ids)
        return logits.to(device1)

class EndpointHandler():

    def __init__(self, model_id = ""):
        model_dir = "pt_fp32"
        model_path = f"{model_dir}/torch_model.pt"
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
        self.split_model = SplitModel(torch.load(model_path))
        self.split_model.eval()
        self.star_line = "***********************************************************"

    def __call__(self, input_data):
        t1 = time.time()
        inputs = input_data.pop("inputs", input_data)
        user_name = inputs["user_name"]
        char_name = inputs["char_name"]
        user_input = inputs["user_input"]
        chats_curled = inputs["chats_curled"]
        while True:
            prompt = template.format(
                char_name = char_name,
                user_name = user_name,
                user_input = "\n".join(user_input)
            )
            input_ids = self.tokenizer(prompt, return_tensors="pt").to("cuda")
            print(f"Token Length: {input_ids.input_ids.size(1)}")
            if input_ids.input_ids.size(1) > 1500:
                chats_curled += 1
                user_input = user_input[chats_curled*2:]
            else: break
        t2 = time.time()
        input_ids = input_ids["input_ids"]
        temperature = 0.5
        max_new_tokens = 50
        with torch.no_grad():
            for _ in range(max_new_tokens):
                attention_mask = torch.ones_like(input_ids).to(device1)
                logits = self.split_model(input_ids, attention_mask)[:, -1] / temperature
                probabilities = torch.softmax(logits, dim=-1)
                sampled_token_ids = torch.multinomial(probabilities, num_samples=1)
                input_ids = torch.cat((input_ids, sampled_token_ids), dim=-1)
                del logits, probabilities, sampled_token_ids
                torch.cuda.empty_cache()
            generated_ids = input_ids.squeeze().tolist()
        t3 = time.time()
        decoded_output = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
        decoded_output = decoded_output.replace(prompt,"").split(f"{user_name}:",1)[0].strip()
        parsed_result = re.sub('\*.*?\*', '', decoded_output).strip()
        if len(parsed_result) != 0: decoded_output = parsed_result
        decoded_output = " ".join(decoded_output.replace("*","").split())
        decoded_output = decoded_output.replace("<USER>", user_name).replace("<BOT>", char_name)
        try:
            parsed_result = decoded_output[:[m.start() for m in re.finditer(r'[.!?]', decoded_output)][-1]+1]
            if len(parsed_result) != 0: decoded_output = parsed_result
        except Exception: pass
        t4 = time.time()
        print(self.star_line)
        print(f"Response: {decoded_output}")
        print(f"Generation Time: {(t3-t2):.2f}")
        print(f"Evaluation Time: {(t4-t1):.2f}")
        print(self.star_line)
        return {
            "message": decoded_output,
            "chats_curled": chats_curled
        }