import torch
import torch.nn.functional as F
import warnings
import sys

warnings.filterwarnings("ignore")


class Chatbot:
    def __init__(self, model_name="stabilityai/stablelm-3b-4e1t", device="cuda"):
        from transformers import (
            AutoTokenizer,
            AutoModelForCausalLM,
            AutoModelForSeq2SeqLM,
        )

        self.model = AutoModelForCausalLM.from_pretrained(
            model_name, torch_dtype=torch.bfloat16, trust_remote_code=True
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

        self.device = device
        self.model.to(device)

    @torch.no_grad()
    def generate_response_greedy(
        self,
        input_text,
        pre_prompt,
        break_word,
        max_length=100,
        temp=0.6,
        name="",
        past_key_vals=None,
        next_id=None,
        verbose=True,
    ):

        if past_key_vals is None:
            inputs = self.tokenizer.encode(
                pre_prompt + input_text + "\n" + name, return_tensors="pt"
            ).to(self.device)
            response_ids = inputs
        else:
            inputs = self.tokenizer.encode(
                input_text + "\n" + name, return_tensors="pt"
            ).to(self.device)
            response_ids = torch.concat((next_id, inputs), dim=-1)
        output = ""
        last_n = ""
        if verbose:
            print(name, end="")
        response_text = ""
        for _ in range(max_length):
            out = self.model.forward(
                input_ids=response_ids, past_key_values=past_key_vals
            )
            next_token_id = torch.multinomial(
                F.softmax(out.logits[:, -1, :] / temp, dim=-1), num_samples=1
            )
            past_key_vals = out.past_key_values
            response_ids = next_token_id
            # print([response_ids[0][-1].to('cpu')])
            output = self.tokenizer.decode([response_ids[0][-1].to("cpu")])
            if verbose:
                print(output, end="")
            response_text += output
            sys.stdout.flush()
            last_n += output
            last_n = last_n[-len(break_word) :]
            if (
                last_n == break_word
                or response_text.strip().endswith("<|endoftext|>")
                or response_text.strip().endswith("[END]")
            ):
                break
        past_kv = past_key_vals
        next_id = response_ids
        return response_text, past_kv, next_id


if __name__ == "__main__":
    from prompts import sales_pre_prompt

    preprompt = sales_pre_prompt
    john = Chatbot()
    break_word = "[USER]"
    name = "[JOHN]"
    log = ""
    past_kv = None
    next_id = None
    print("type: exit, quit or stop to end the chat")
    print("Chat started:")
    while True:
        user_input = input(" ")
        if user_input.lower() in ["exit", "quit", "stop"]:
            break

        response, past_kv, next_id = john.generate_response_greedy(
            " " + user_input,
            preprompt + log,
            break_word,
            max_length=100000,
            name=name,
            past_key_vals=past_kv,
            next_id=next_id,
            verbose=True,
            temp=0.2,
        )

        log += " " + user_input + "\n" + name + response
        print("------\n" + log + "\n----------")
